chore: import upstream snapshot with attribution

tensorflow/python/framework/experimental/BUILD

@@ -0,0 +1,238 @@
+# Experimental Unified APIs for Eager and Graph modes.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_python_pybind_extension")
+load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
+load(
+    "//tensorflow/tools/test:performance.bzl",
+    "cuda_py_benchmark_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_python_pybind_extension(
+    name = "_unified_api",
+    srcs = ["unified_api.cc"],
+    enable_stub_generation = True,
+    features = ["-layering_check"],
+    pytype_srcs = [
+        "_unified_api.pyi",
+    ],
+    starlark_only = True,
+    visibility = [
+        "//tensorflow/python:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/lib/llvm_rtti",
+        "//tensorflow/core/platform:refcount",
+        "//tensorflow/python:unified_api_pywrap_required_headers",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@com_google_absl//absl/types:span",
+        "@pybind11",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "_tape",
+    srcs = ["tape.cc"],
+    features = ["-layering_check"],
+    starlark_only = True,
+    visibility = [
+        "//tensorflow/python:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/lib/llvm_rtti",
+        "//tensorflow/python:unified_api_pywrap_required_headers",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
+        "@pybind11",
+    ] + if_pywrap(
+        if_true = [
+            "//tensorflow/c/experimental/gradients/tape:tape_context",
+            "//tensorflow/c/experimental/gradients:math_grad",
+            "//tensorflow/c/experimental/gradients:nn_grad",
+        ],
+    ),
+)
+
+tf_python_pybind_extension(
+    name = "_math_ops",
+    srcs = ["math_ops.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_math_ops.pyi",
+    ],
+    starlark_only = True,
+    visibility = [
+        "//tensorflow/python:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/lib/llvm_rtti",
+        "//tensorflow/python:unified_api_pywrap_required_headers",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@com_google_absl//absl/types:span",
+        "@pybind11",
+    ] + if_pywrap(
+        if_true = [
+            "//tensorflow/c/experimental/ops:math_ops",
+        ],
+    ),
+)
+
+tf_python_pybind_extension(
+    name = "_nn_ops",
+    srcs = ["nn_ops.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_nn_ops.pyi",
+    ],
+    starlark_only = True,
+    visibility = [
+        "//tensorflow/python:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/lib/llvm_rtti",
+        "//tensorflow/python:unified_api_pywrap_required_headers",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@com_google_absl//absl/types:span",
+        "@pybind11",
+    ] + if_pywrap(
+        if_true = [
+            "//tensorflow/c/experimental/ops:nn_ops",
+        ],
+    ),
+)
+
+py_library(
+    name = "gradient_registry",
+    srcs = ["gradient_registry.py"],
+    strict_deps = True,
+    deps = [":_tape"],
+)
+
+py_library(
+    name = "math_ops",
+    srcs = ["math_ops.py"],
+    strict_deps = True,
+    deps = [
+        ":_math_ops",
+        ":context_stack",
+    ],
+)
+
+py_library(
+    name = "nn_ops",
+    srcs = ["nn_ops.py"],
+    strict_deps = True,
+    deps = [
+        ":_nn_ops",
+        ":context_stack",
+    ],
+)
+
+py_library(
+    name = "tape",
+    srcs = ["tape.py"],
+    strict_deps = True,
+    deps = [
+        ":_tape",
+        ":context_stack",
+        ":gradient_registry",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_library(
+    name = "def_function",
+    srcs = ["def_function.py"],
+    strict_deps = True,
+    deps = [
+        ":_unified_api",
+        ":context_stack",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_library(
+    name = "thread_local_stack",
+    srcs = ["thread_local_stack.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "context_stack",
+    srcs = ["context_stack.py"],
+    strict_deps = True,
+    deps = [":thread_local_stack"],
+)
+
+cuda_py_strict_test(
+    name = "unified_api_test",
+    size = "small",
+    srcs = ["unified_api_test.py"],
+    tags = [
+        # Note(srbs): These python bindings are not
+        # exported as part of the pip package yet so
+        # this test is disabled.
+        "no_pip",
+        "no_windows",  # b/168218876
+    ],
+    deps = [
+        ":_unified_api",
+        ":context_stack",
+        ":def_function",
+        ":math_ops",
+        ":nn_ops",
+        ":tape",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_grad",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_benchmark_test(
+    name = "graph_building_test",
+    size = "small",
+    srcs = ["graph_building_test.py"],
+    deps = [
+        "//tensorflow/core/config:flags_py",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:resource_variable_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/framework/experimental/_math_ops.pyi

@@ -0,0 +1,22 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def add(*args, **kwargs): ...
+def div_no_nan(*args, **kwargs): ...
+def log1p(*args, **kwargs): ...
+def mat_mul(*args, **kwargs): ...
+def mul(*args, **kwargs): ...
+def neg(*args, **kwargs): ...
+def sub(*args, **kwargs): ...

tensorflow/python/framework/experimental/_nn_ops.pyi

@@ -0,0 +1,17 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def relu(*args, **kwargs): ...
+def sparse_softmax_cross_entropy_with_logits(*args, **kwargs): ...

tensorflow/python/framework/experimental/_unified_api.pyi

@@ -0,0 +1,52 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class AbstractContext:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def CreateOperation(self, *args, **kwargs): ...
+    def RegisterFunction(self, arg0) -> None: ...
+    def RemoveFunction(self, arg0: str) -> None: ...
+
+class AbstractFunction:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class AbstractOperation:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def AddInput(self, arg0) -> None: ...
+    def DeviceName(self) -> str: ...
+    def Execute(self, *args, **kwargs): ...
+    def Name(self) -> str: ...
+    def Reset(self, arg0: str, arg1: str) -> None: ...
+    def SetAttrType(self, arg0: str, arg1) -> None: ...
+    def SetDeviceName(self, arg0: str) -> None: ...
+    def SetOpName(self, arg0: str) -> None: ...
+
+class AbstractTensorHandle:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def DataType(self, *args, **kwargs): ...
+    def numpy(self) -> object: ...
+
+class ImmediateExecutionContext(AbstractContext):
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TracingContext(AbstractContext):
+    def __init__(self, *args, **kwargs) -> None: ...
+    def AddParameter(self, *args, **kwargs): ...
+    def Finalize(self, *args, **kwargs): ...
+
+def EagerContextToImmediateExecutionContext(*args, **kwargs): ...
+def EagerTensorToImmediateExecutionTensorHandle(arg0: object) -> AbstractTensorHandle: ...
+def NewTracingContext(*args, **kwargs): ...
+def SetTracingImplementation(arg0: str) -> None: ...

tensorflow/python/framework/experimental/context_stack.py

@@ -0,0 +1,36 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Thread-local context manager stack."""
+
+import contextlib
+
+from tensorflow.python.framework.experimental import thread_local_stack
+
+_default_ctx_stack = thread_local_stack.ThreadLocalStack()
+
+
+def get_default():
+  """Returns the default execution context."""
+  return _default_ctx_stack.peek()
+
+
+@contextlib.contextmanager
+def set_default(ctx):
+  """Returns a contextmanager with `ctx` as the default execution context."""
+  try:
+    _default_ctx_stack.push(ctx)
+    yield
+  finally:
+    _default_ctx_stack.pop()

tensorflow/python/framework/experimental/def_function.py

@@ -0,0 +1,70 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Experimental impl of tf.function using unified APIs, for testing only."""
+
+from tensorflow.python.framework.experimental import _unified_api
+from tensorflow.python.framework.experimental import context_stack as context_lib
+from tensorflow.python.util import nest
+
+NewTracingContext = _unified_api.NewTracingContext
+
+
+class Function(object):
+  """Helper for tf.function."""
+
+  def __init__(self, func, name=None):
+    self._python_func = func
+    # TODO(srbs): Uniquify this name.
+    self.name = name or func.__name__
+
+  def __call__(self, *args, **kwargs):
+    # Flatten arguments.
+    flat_args = nest.flatten(args, expand_composites=True)
+    flat_kwargs = nest.flatten(kwargs, expand_composites=True)
+    all_args = flat_args + flat_kwargs
+
+    # Trace
+    outer_ctx = context_lib.get_default()
+    ctx = NewTracingContext(self.name)
+    with context_lib.set_default(ctx):
+      # TODO(srbs): Iterating over list of inputs is a known performance
+      # bottleneck. Add a pybind API for this.
+      inputs = [ctx.AddParameter(arg.DataType()) for arg in all_args]
+      structured_args = nest.pack_sequence_as(args, inputs[:len(flat_args)])
+      structured_kwargs = nest.pack_sequence_as(kwargs, inputs[len(flat_args):])
+      structured_outputs = self._python_func(*structured_args,
+                                             **structured_kwargs)
+
+      py_outputs = nest.flatten(structured_outputs, expand_composites=True)
+      num_outputs = len(py_outputs)
+      # TODO(srbs): Drop Nones before calling Finalize.
+      finalized_f = ctx.Finalize(py_outputs)
+      outer_ctx.RegisterFunction(finalized_f)
+
+    # Build call op
+    call_op = outer_ctx.CreateOperation(self.name, "")
+    call_op.SetOpName(self.name)
+    for arg in all_args:
+      call_op.AddInput(arg)
+    call_op_outputs = call_op.Execute(num_outputs)
+
+    # Cleanup
+    outer_ctx.RemoveFunction(self.name)
+
+    return nest.pack_sequence_as(structured_outputs, call_op_outputs)
+
+
+def function(func):
+  return Function(func)

tensorflow/python/framework/experimental/gradient_registry.py

@@ -0,0 +1,23 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Global GradientRegistry."""
+
+from tensorflow.python.framework.experimental import _tape
+
+_GRADIENT_REGISTRY_GLOBAL = _tape.GradientRegistry()
+
+
+def get_global_registry():
+  return _GRADIENT_REGISTRY_GLOBAL

tensorflow/python/framework/experimental/graph_building_test.py

@@ -0,0 +1,89 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for adding ops to a graph."""
+
+import timeit
+
+from tensorflow.core.config import flags
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import func_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.platform import test
+
+
+@test_util.add_graph_building_optimization_tests
+class GraphBuildingBenchmark(test.Benchmark):
+
+  def _computeAddOpDuration(self, num_ops, num_iters):
+    def add_op_to_graph(num_ops):
+      with func_graph.FuncGraph("add").as_default():
+        a = gen_array_ops.placeholder(dtypes.float32)
+        b = gen_array_ops.placeholder(dtypes.float32)
+        for _ in range(num_ops):
+          gen_math_ops.add(a, b)
+
+    runtimes = timeit.repeat(
+        lambda: add_op_to_graph(num_ops), repeat=10, number=num_iters)
+    return min(runtimes) / num_iters
+
+  def _computeReadVariableOpDuration(self, num_ops, num_iters):
+    def add_op_to_graph(num_ops):
+      with func_graph.FuncGraph("resource").as_default():
+        handle = resource_variable_ops.var_handle_op(
+            dtype=dtypes.int32, shape=[])
+        resource_variable_ops.assign_variable_op(
+            handle, constant_op.constant(1, dtype=dtypes.int32))
+        for _ in range(num_ops):
+          gen_resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32)
+
+    runtimes = timeit.repeat(
+        lambda: add_op_to_graph(num_ops), repeat=10, number=num_iters)
+    return min(runtimes) / num_iters
+
+  def benchmarkAddOp(self):
+    num_ops = 100
+    num_iters = 10
+    duration = self._computeAddOpDuration(num_ops, num_iters)
+    name = "BenchmarkAddOp"
+    if flags.config().graph_building_optimization.value():
+      name += "WithGraphBuildingOptimization"
+    self.report_benchmark(
+        name=name,
+        iters=num_iters,
+        wall_time=duration,
+        extras={"num_ops": num_ops})
+
+  def benchmarkResourceVariableOp(self):
+    num_ops = 100
+    num_iters = 10
+    duration = self._computeReadVariableOpDuration(num_ops, num_iters)
+    name = "BenchmarkReadVariableOp"
+    if flags.config().graph_building_optimization.value():
+      name += "WithGraphBuildingOptimization"
+    self.report_benchmark(
+        name=name,
+        iters=num_iters,
+        wall_time=duration,
+        extras={"num_ops": num_ops})
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/framework/experimental/math_ops.cc

@@ -0,0 +1,96 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/experimental/ops/math_ops.h"
+
+#include <pybind11/stl.h>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+using tensorflow::AbstractContext;
+using tensorflow::AbstractTensorHandle;
+
+namespace tensorflow {
+PYBIND11_MODULE(_math_ops, m) {
+  m.def("add", [](AbstractContext* ctx, AbstractTensorHandle* a,
+                  AbstractTensorHandle* b, const char* name) {
+    AbstractTensorHandle* output;
+    if (!name) {
+      name = "Add";
+    }
+    MaybeRaiseRegisteredFromStatus(ops::AddV2(ctx, a, b, &output, name));
+    return output;
+  });
+  m.def("mat_mul", [](AbstractContext* ctx, AbstractTensorHandle* a,
+                      AbstractTensorHandle* b, const char* name) {
+    AbstractTensorHandle* output;
+    if (!name) {
+      name = "MatMul";
+    }
+    MaybeRaiseRegisteredFromStatus(ops::MatMul(ctx, a, b, &output,
+                                               /*transpose_a=*/false,
+                                               /*transpose_b=*/false, name));
+    return output;
+  });
+  m.def("neg",
+        [](AbstractContext* ctx, AbstractTensorHandle* a, const char* name) {
+          AbstractTensorHandle* output;
+          if (!name) {
+            name = "Neg";
+          }
+          MaybeRaiseRegisteredFromStatus(ops::Neg(ctx, a, &output, name));
+          return output;
+        });
+  m.def("sub", [](AbstractContext* ctx, AbstractTensorHandle* a,
+                  AbstractTensorHandle* b, const char* name) {
+    AbstractTensorHandle* output;
+    if (!name) {
+      name = "Sub";
+    }
+    MaybeRaiseRegisteredFromStatus(ops::Sub(ctx, a, b, &output, name));
+    return output;
+  });
+  m.def("mul", [](AbstractContext* ctx, AbstractTensorHandle* a,
+                  AbstractTensorHandle* b, const char* name) {
+    AbstractTensorHandle* output;
+    if (!name) {
+      name = "Mul";
+    }
+    MaybeRaiseRegisteredFromStatus(ops::Mul(ctx, a, b, &output, name));
+    return output;
+  });
+  m.def("log1p",
+        [](AbstractContext* ctx, AbstractTensorHandle* a, const char* name) {
+          AbstractTensorHandle* output;
+          if (!name) {
+            name = "Log1p";
+          }
+          MaybeRaiseRegisteredFromStatus(ops::Log1p(ctx, a, &output, name));
+          return output;
+        });
+  m.def("div_no_nan", [](AbstractContext* ctx, AbstractTensorHandle* a,
+                         AbstractTensorHandle* b, const char* name) {
+    AbstractTensorHandle* output;
+    if (!name) {
+      name = "DivNoNan";
+    }
+    MaybeRaiseRegisteredFromStatus(ops::DivNoNan(ctx, a, b, &output, name));
+    return output;
+  });
+}
+}  // namespace tensorflow

tensorflow/python/framework/experimental/math_ops.py

@@ -0,0 +1,53 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Experimental impl for gen_math_ops.py using unified APIs, for testing."""
+
+from tensorflow.python.framework.experimental import _math_ops
+from tensorflow.python.framework.experimental import context_stack as context
+
+
+def add(a, b, name=None):
+  ctx = context.get_default()
+  return _math_ops.add(ctx, a, b, name)
+
+
+def mat_mul(a, b, name=None):
+  ctx = context.get_default()
+  return _math_ops.mat_mul(ctx, a, b, name)
+
+
+def neg(a, name=None):
+  ctx = context.get_default()
+  return _math_ops.neg(ctx, a, name)
+
+
+def sub(a, b, name=None):
+  ctx = context.get_default()
+  return _math_ops.sub(ctx, a, b, name)
+
+
+def mul(a, b, name=None):
+  ctx = context.get_default()
+  return _math_ops.mul(ctx, a, b, name)
+
+
+def log1p(a, name=None):
+  ctx = context.get_default()
+  return _math_ops.log1p(ctx, a, name)
+
+
+def div_no_nan(a, b, name=None):
+  ctx = context.get_default()
+  return _math_ops.div_no_nan(ctx, a, b, name)

tensorflow/python/framework/experimental/nn_ops.cc

@@ -0,0 +1,77 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/c/experimental/ops/nn_ops.h"
+
+#include <pybind11/stl.h>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+using tensorflow::AbstractContext;
+using tensorflow::AbstractTensorHandle;
+
+namespace tensorflow {
+PYBIND11_MODULE(_nn_ops, m) {
+  m.def("relu",
+        [](AbstractContext* ctx, AbstractTensorHandle* a, const char* name) {
+          if (ctx == nullptr) {
+            throw pybind11::type_error("relu() argument 'ctx' cannot be None");
+          }
+          if (a == nullptr) {
+            throw pybind11::type_error("relu() argument 'a' cannot be None");
+          }
+          AbstractTensorHandle* output;
+          if (!name) {
+            name = "Relu";
+          }
+          MaybeRaiseRegisteredFromStatus(ops::Relu(ctx, a, &output, name));
+          return output;
+        });
+
+  m.def(
+      "sparse_softmax_cross_entropy_with_logits",
+      [](AbstractContext* ctx, AbstractTensorHandle* features,
+         AbstractTensorHandle* labels, const char* name) {
+        if (ctx == nullptr) {
+          throw pybind11::type_error(
+              "sparse_softmax_cross_entropy_with_logits() argument 'ctx' "
+              "cannot be None");
+        }
+        if (features == nullptr) {
+          throw pybind11::type_error(
+              "sparse_softmax_cross_entropy_with_logits() argument 'features' "
+              "cannot be None");
+        }
+        if (labels == nullptr) {
+          throw pybind11::type_error(
+              "sparse_softmax_cross_entropy_with_logits() argument 'labels' "
+              "cannot be None");
+        }
+        AbstractTensorHandle* loss;
+        AbstractTensorHandle* backprop;
+        if (!name) {
+          name = "SparseSoftmaxCrossEntropyWithLogits";
+        }
+        absl::Status status = ops::SparseSoftmaxCrossEntropyWithLogits(
+            ctx, features, labels, &loss, &backprop, name);
+        MaybeRaiseRegisteredFromStatus(status);
+        backprop->Unref();
+        return loss;
+      });
+}
+}  // namespace tensorflow

tensorflow/python/framework/experimental/nn_ops.py

@@ -0,0 +1,29 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Experimental impl for gen_nn_ops.py using unified APIs, for testing only."""
+
+from tensorflow.python.framework.experimental import _nn_ops
+from tensorflow.python.framework.experimental import context_stack as context
+
+
+def relu(a, name=None):
+  ctx = context.get_default()
+  return _nn_ops.relu(ctx, a, name)
+
+
+def sparse_softmax_cross_entropy_with_logits(logits, labels, name=None):
+  ctx = context.get_default()
+  return _nn_ops.sparse_softmax_cross_entropy_with_logits(
+      ctx, logits, labels, name)

tensorflow/python/framework/experimental/tape.cc

@@ -0,0 +1,79 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <pybind11/stl.h>
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/c/eager/gradients.h"
+#include "tensorflow/c/experimental/gradients/math_grad.h"
+#include "tensorflow/c/experimental/gradients/nn_grad.h"
+#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+namespace py = pybind11;
+
+namespace tensorflow {
+namespace gradients {
+
+absl::Status RegisterGradients(GradientRegistry* registry) {
+  // TODO(srbs): Rename ops::Add and AddRegisterer to AddV2.
+  TF_RETURN_IF_ERROR(registry->Register("AddV2", AddRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("Exp", ExpRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("MatMul", MatMulRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("Relu", ReluRegisterer));
+  TF_RETURN_IF_ERROR(
+      registry->Register("SparseSoftmaxCrossEntropyWithLogits",
+                         SparseSoftmaxCrossEntropyWithLogitsRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("Neg", NegRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("Sub", SubRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("Mul", MulRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("Log1p", Log1pRegisterer));
+  TF_RETURN_IF_ERROR(registry->Register("DivNoNan", DivNoNanRegisterer));
+  return absl::OkStatus();
+}
+
+PYBIND11_MODULE(_tape, m) {
+  py::class_<Tape>(m, "Tape")
+      .def(py::init([](bool persistent) { return new Tape(persistent); }))
+      .def("Watch", [](Tape* self, AbstractTensorHandle* t) { self->Watch(t); })
+      .def("ComputeGradient",
+           [](Tape* self, AbstractContext* ctx,
+              std::vector<AbstractTensorHandle*> target_tensors,
+              std::vector<AbstractTensorHandle*> source_tensors,
+              std::vector<AbstractTensorHandle*> output_gradients) {
+             std::vector<AbstractTensorHandle*> results(source_tensors.size());
+             absl::Status s = self->ComputeGradient(
+                 ctx, target_tensors, source_tensors, output_gradients,
+                 absl::MakeSpan(results));
+             MaybeRaiseRegisteredFromStatus(s);
+             return results;
+           });
+  py::class_<GradientRegistry>(m, "GradientRegistry").def(py::init([]() {
+    auto registry = new GradientRegistry();
+    MaybeRaiseRegisteredFromStatus(RegisterGradients(registry));
+    return registry;
+  }));
+  py::class_<TapeContext, AbstractContext>(m, "TapeContext")
+      .def(py::init(
+          [](AbstractContext* ctx, Tape* tape, GradientRegistry* registry) {
+            return new TapeContext(ctx, tape, *registry);
+          }));
+}
+}  // namespace gradients
+}  // namespace tensorflow

tensorflow/python/framework/experimental/tape.py

@@ -0,0 +1,53 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Experimental impl for GradientTape using unified APIs, for testing only."""
+
+from tensorflow.python.framework.experimental import _tape
+from tensorflow.python.framework.experimental import context_stack
+from tensorflow.python.framework.experimental import gradient_registry
+from tensorflow.python.util import nest
+
+
+class GradientTape(object):
+  """GradientTape using the unified API."""
+
+  def __init__(self, persistent=False):
+    self._c_tape = _tape.Tape(persistent)
+    ctx = context_stack.get_default()
+    self._tape_context = _tape.TapeContext(
+        ctx, self._c_tape, gradient_registry.get_global_registry())
+    self._ctx_manager = None
+
+  def watch(self, t):
+    self._c_tape.Watch(t)
+
+  # TODO(srbs): Add support for unconnected_gradients.
+  def gradient(self, targets, sources, output_gradients=None):
+    ctx = context_stack.get_default()
+    flat_targets = nest.flatten(targets)
+    flat_sources = nest.flatten(sources)
+    out_grads = self._c_tape.ComputeGradient(ctx, flat_targets, flat_sources,
+                                             output_gradients or [])
+    return nest.pack_sequence_as(sources, out_grads)
+
+  def __enter__(self):
+    """Enters a context inside which operations are recorded on this tape."""
+    self._ctx_manager = context_stack.set_default(self._tape_context)
+    self._ctx_manager.__enter__()
+    return self
+
+  def __exit__(self, typ, value, traceback):
+    self._ctx_manager.__exit__(typ, value, traceback)
+    self._ctx_manager = None

tensorflow/python/framework/experimental/thread_local_stack.py

@@ -0,0 +1,35 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Thread-local stack."""
+
+import threading
+
+
+# TODO(srbs): Move this to C++.
+class ThreadLocalStack(threading.local):
+  """A thread-local stack of objects for providing implicit defaults."""
+
+  def __init__(self):
+    super(ThreadLocalStack, self).__init__()
+    self._stack = []
+
+  def peek(self):
+    return self._stack[-1] if self._stack else None
+
+  def push(self, ctx):
+    return self._stack.append(ctx)
+
+  def pop(self):
+    self._stack.pop()

tensorflow/python/framework/experimental/unified_api.cc

@@ -0,0 +1,267 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <pybind11/stl.h>
+
+#include <memory>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/c/eager/abstract_context.h"
+#include "tensorflow/c/eager/abstract_function.h"
+#include "tensorflow/c/eager/abstract_operation.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_internal.h"
+#include "tensorflow/c/eager/c_api_unified_experimental.h"
+#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/eager/tfe_context_internal.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/c/safe_ptr.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/python/eager/pywrap_tensor.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+
+namespace py = pybind11;
+
+using tensorflow::AbstractContext;
+using tensorflow::AbstractContextPtr;
+using tensorflow::AbstractFunction;
+using tensorflow::AbstractOperation;
+using tensorflow::AbstractOperationPtr;
+using tensorflow::AbstractTensorHandle;
+using tensorflow::AbstractTensorHandlePtr;
+using tensorflow::OutputList;
+
+using tensorflow::tracing::TracingContext;
+using tensorflow::tracing::TracingOperation;
+using tensorflow::tracing::TracingTensorHandle;
+
+using tensorflow::ImmediateContextPtr;
+using tensorflow::ImmediateExecutionContext;
+using tensorflow::ImmediateExecutionTensorHandle;
+
+using tensorflow::dyn_cast;
+using tensorflow::isa;
+using tensorflow::unwrap;
+using tensorflow::wrap;
+
+using tensorflow::DataType;
+using tensorflow::make_safe;
+using tensorflow::MaybeRaiseRegisteredFromStatus;
+using tensorflow::MaybeRaiseRegisteredFromTFStatus;
+using tensorflow::Pyo;
+using tensorflow::Safe_TF_StatusPtr;
+using tensorflow::Status;
+using tensorflow::string;
+using tensorflow::TFE_TensorHandleToNumpy;
+
+using tensorflow::errors::Internal;
+using tensorflow::errors::InvalidArgument;
+
+PYBIND11_MODULE(_unified_api, m) {
+  // Context creation functions.
+  m.def("SetTracingImplementation", [](const char* impl) {
+    Safe_TF_StatusPtr status = make_safe(TF_NewStatus());
+    TF_SetTracingImplementation(impl, status.get());
+    MaybeRaiseRegisteredFromStatus(status->status);
+  });
+  m.def("NewTracingContext", [](const char* fn_name) {
+    Safe_TF_StatusPtr status = make_safe(TF_NewStatus());
+    auto* ctx = unwrap(TF_CreateFunction(fn_name, status.get()));
+    MaybeRaiseRegisteredFromTFStatus(status.get());
+    if (!ctx) {
+      MaybeRaiseRegisteredFromStatus(
+          absl::InternalError("TF_CreateFunction returned nullptr"));
+    }
+    if (!isa<TracingContext>(ctx)) {
+      // TODO(srbs): Add a helper to convert the kind enum to a user-friendly
+      // string.
+      MaybeRaiseRegisteredFromStatus(absl::InternalError(
+          absl::StrCat("TF_CreateFunction must return a TracingContext, found ",
+                       ctx->getKind())));
+    }
+    return dyn_cast<TracingContext>(ctx);
+  });
+  m.def("EagerContextToImmediateExecutionContext", [](py::handle& obj) {
+    TFE_Context* ctx = static_cast<TFE_Context*>(
+        PyCapsule_GetPointer(obj.ptr(), "TFE_Context"));
+    if (!ctx) {
+      MaybeRaiseRegisteredFromStatus(
+          absl::InvalidArgumentError("TFE_Context is nullptr"));
+    }
+    return unwrap(ctx);
+  });
+
+  // Unified execution context.
+  py::class_<AbstractContext, AbstractContextPtr>(m, "AbstractContext")
+      .def("CreateOperation",
+           [](AbstractContext* self, const char* op,
+              const char* raw_device_name) {
+             auto operation = self->CreateOperation();
+             (void)operation->Reset(op, raw_device_name);
+             return operation;
+           })
+      .def("RegisterFunction",
+           [](AbstractContext* self, AbstractFunction* f) {
+             Status s = self->RegisterFunction(f);
+             MaybeRaiseRegisteredFromStatus(s);
+           })
+      .def("RemoveFunction", [](AbstractContext* self, const string& func) {
+        Status s = self->RemoveFunction(func);
+        MaybeRaiseRegisteredFromStatus(s);
+      });
+
+  py::class_<TracingContext, AbstractContext>(m, "TracingContext")
+      .def("AddParameter",
+           [](TracingContext* self, DataType dtype) {
+             TracingTensorHandle* handle = nullptr;
+             // TODO(srbs): Add shape argument to this function.
+             tensorflow::PartialTensorShape shape;
+             Status s = self->AddParameter(dtype, shape, &handle);
+             MaybeRaiseRegisteredFromStatus(s);
+             return static_cast<AbstractTensorHandle*>(handle);
+           })
+      .def("Finalize", [](TracingContext* self, py::handle& outputs) {
+        // TODO(srbs): Using OutputList seems like an overkill here. Should we
+        // simply pass in an absl::Span?
+        OutputList output_list;
+        if (outputs.ptr() != Py_None) {
+          if (!PyList_Check(outputs.ptr())) {
+            MaybeRaiseRegisteredFromStatus(absl::InvalidArgumentError(
+                "must provide a list of Tensors as inputs"));
+          }
+          Py_ssize_t len = PyList_Size(outputs.ptr());
+          output_list.outputs.resize(len);
+          for (Py_ssize_t i = 0; i < len; ++i) {
+            PyObject* elem = PyList_GetItem(outputs.ptr(), i);
+            if (!elem) {
+              MaybeRaiseRegisteredFromStatus(absl::InvalidArgumentError(
+                  absl::StrCat("Tensor at index  ", i, " is None.")));
+            }
+            py::handle elem_h = elem;
+            AbstractTensorHandle* handle = elem_h.cast<AbstractTensorHandle*>();
+            if (!isa<TracingTensorHandle>(handle)) {
+              MaybeRaiseRegisteredFromStatus(
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Tensor at index  ", i, " is not a graph tensor.")));
+            }
+            output_list.outputs[i] = handle;
+          }
+        }
+        AbstractFunction* f = nullptr;
+        Status s = self->Finalize(&output_list, &f);
+        MaybeRaiseRegisteredFromStatus(s);
+        return f;
+      });
+
+  // Note: This does not take ownership of the C++ context, the lifetime of
+  // which is managed by the python `Context` and is expected to outlive this
+  // object.
+  // TODO(srbs): Make AbstractContext refcounted so that the above comment is
+  // not needed.
+  py::class_<ImmediateExecutionContext, AbstractContext,
+             std::unique_ptr<ImmediateExecutionContext, py::nodelete>>
+      ImmediateExecutionContext(m, "ImmediateExecutionContext");
+
+  // Unified execution operation.
+  py::class_<AbstractOperation, AbstractOperationPtr>(m, "AbstractOperation")
+      .def("Reset",
+           [](AbstractOperation* self, const char* op,
+              const char* raw_device_name) {
+             Status s = self->Reset(op, raw_device_name);
+             MaybeRaiseRegisteredFromStatus(s);
+           })
+      .def("SetOpName",
+           [](AbstractOperation* self, const char* op_name) {
+             // TODO(srbs): We could provide SetOpName on TracingOperation
+             // but then we need to do a hasattr check or try/pass in python.
+             if (isa<TracingOperation>(self)) {
+               auto tracing_op = reinterpret_cast<TracingOperation*>(self);
+               Status s = tracing_op->SetOpName(op_name);
+               MaybeRaiseRegisteredFromStatus(s);
+             }
+           })
+      .def("Name", &AbstractOperation::Name)
+      .def("DeviceName", &AbstractOperation::DeviceName)
+      .def("SetDeviceName",
+           [](AbstractOperation* self, const char* name) {
+             Status s = self->SetDeviceName(name);
+             MaybeRaiseRegisteredFromStatus(s);
+           })
+      .def("AddInput",
+           [](AbstractOperation* self, AbstractTensorHandle* input) {
+             Status s = self->AddInput(input);
+             MaybeRaiseRegisteredFromStatus(s);
+           })
+      .def("SetAttrType",
+           [](AbstractOperation* self, const char* attr_name, DataType value) {
+             Status s = self->SetAttrType(attr_name, value);
+             MaybeRaiseRegisteredFromStatus(s);
+           })
+      .def("Execute", [](AbstractOperation* self, int num_outputs) {
+        std::vector<AbstractTensorHandle*> outputs(num_outputs);
+        MaybeRaiseRegisteredFromStatus(
+            self->Execute(absl::MakeSpan(outputs), &num_outputs));
+        return outputs;
+      });
+
+  // Unified execution tensor handle.
+  py::class_<AbstractTensorHandle, AbstractTensorHandlePtr>(
+      m, "AbstractTensorHandle")
+      .def("DataType", &AbstractTensorHandle::DataType)
+      .def("numpy", [](AbstractTensorHandle* self) {
+        // TODO(srbs): Export this on ImmediateExecutionTensorHandle only.
+        if (!isa<ImmediateExecutionTensorHandle>(self)) {
+          // TODO(srbs): Add a helper to convert the kind enum to a
+          // user-friendly string.
+          MaybeRaiseRegisteredFromStatus(absl::InternalError(absl::StrCat(
+              "AbstractTensorHandle.numpy() must be called with an ",
+              "ImmediateExecutionTensorHandle found type: ", self->getKind())));
+        }
+        TF_Status s;
+        TFE_TensorHandle* handle =
+            wrap(dyn_cast<ImmediateExecutionTensorHandle>(self));
+        auto result = TFE_TensorHandleToNumpy(handle, &s);
+        MaybeRaiseRegisteredFromStatus(s.status);
+        return Pyo(result);
+      });
+
+  m.def("EagerTensorToImmediateExecutionTensorHandle", [](py::object handle) {
+    if (!EagerTensor_CheckExact(handle.ptr())) {
+      MaybeRaiseRegisteredFromStatus(absl::InvalidArgumentError(
+          "EagerTensorToImmediateExecutionTensorHandle called "
+          "with non-EagerTensor."));
+    }
+    TFE_TensorHandle* eager_tensor = EagerTensor_Handle(handle.ptr());
+    auto t = static_cast<AbstractTensorHandle*>(unwrap(eager_tensor));
+    t->Ref();
+    return t;
+  });
+
+  py::class_<AbstractFunction,
+             std::unique_ptr<AbstractFunction, tsl::core::RefCountDeleter>>
+      AbstractFunction(m, "AbstractFunction");
+}

tensorflow/python/framework/experimental/unified_api_test.py

@@ -0,0 +1,505 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Unified APIs' python bindings."""
+
+import timeit
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework.experimental import _unified_api
+from tensorflow.python.framework.experimental import context_stack as context_lib
+from tensorflow.python.framework.experimental import def_function
+from tensorflow.python.framework.experimental import math_ops as unified_math_ops
+from tensorflow.python.framework.experimental import nn_ops as unified_nn_ops
+from tensorflow.python.framework.experimental import tape as tape_lib
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import test
+
+SetTracingImplementation = _unified_api.SetTracingImplementation
+TensorCastHelper = _unified_api.EagerTensorToImmediateExecutionTensorHandle
+
+
+def get_immediate_execution_context():
+  context._reset_context()
+  context.context().ensure_initialized()
+  return _unified_api.EagerContextToImmediateExecutionContext(
+      context.context()._handle)
+
+
+def maybe_cast(t, perform_cast):
+  if perform_cast:
+    return TensorCastHelper(t)
+  return t
+
+
+class UnifiedApiTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testAdd(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      return unified_math_ops.add(a, b)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1., 2.]))
+      b = TensorCastHelper(constant_op.constant([3., 4.]))
+
+      func_output = def_function.function(model)(a, b)
+      self.assertAllEqual(func_output.numpy(), [4., 6.])
+
+      eager_output = model(a, b)
+      self.assertAllEqual(eager_output.numpy(), [4., 6.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testAddGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(a)
+        tape.watch(b)
+        result = unified_math_ops.add(a, b)
+      grads = tape.gradient(result, [a, b])
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1., 2.]))
+      b = TensorCastHelper(constant_op.constant([3., 4.]))
+
+      func_outputs = def_function.function(model)(a, b)
+      self.assertAllEqual(func_outputs[0].numpy(), [1.0, 1.0])
+      self.assertAllEqual(func_outputs[1].numpy(), [1.0, 1.0])
+
+      eager_outputs = model(a, b)
+      self.assertAllEqual(eager_outputs[0].numpy(), [1.0, 1.0])
+      self.assertAllEqual(eager_outputs[1].numpy(), [1.0, 1.0])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testRelu(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(t):
+      return unified_nn_ops.relu(t)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      positive = TensorCastHelper(constant_op.constant([1.]))
+      negative = TensorCastHelper(constant_op.constant([-1.]))
+
+      model_fn = def_function.function(model)
+      func_output = model_fn(positive)
+      self.assertAllEqual(func_output.numpy(), [1.])
+      func_output = model_fn(negative)
+      self.assertAllEqual(func_output.numpy(), [0.])
+
+      eager_output = model(positive)
+      self.assertAllEqual(eager_output.numpy(), [1.])
+      eager_output = model(negative)
+      self.assertAllEqual(eager_output.numpy(), [0.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testReluGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(t):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(t)
+        result = unified_nn_ops.relu(t)
+      grads = tape.gradient(result, t)
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      positive = TensorCastHelper(constant_op.constant([1.]))
+      negative = TensorCastHelper(constant_op.constant([-1.]))
+
+      model_fn = def_function.function(model)
+      func_output = model_fn(positive)
+      self.assertAllEqual(func_output.numpy(), [1.])
+      func_output = model_fn(negative)
+      self.assertAllEqual(func_output.numpy(), [0.])
+
+      eager_output = model(positive)
+      self.assertAllEqual(eager_output.numpy(), [1.])
+      eager_output = model(negative)
+      self.assertAllEqual(eager_output.numpy(), [0.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testNeg(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a):
+      return unified_math_ops.neg(a)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([2.]))
+
+      func_output = def_function.function(model)(a)
+      self.assertAllEqual(func_output.numpy(), [-2.])
+
+      eager_output = model(a)
+      self.assertAllEqual(eager_output.numpy(), [-2.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testNegGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(a)
+        result = unified_math_ops.neg(a)
+      grads = tape.gradient(result, a)
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([2.]))
+
+      func_outputs = def_function.function(model)(a)
+      self.assertAllEqual(func_outputs.numpy(), [-1.0])
+
+      eager_outputs = model(a)
+      self.assertAllEqual(eager_outputs.numpy(), [-1.0])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testSub(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      return unified_math_ops.sub(a, b)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1., 2.]))
+      b = TensorCastHelper(constant_op.constant([3., 4.]))
+
+      func_output = def_function.function(model)(a, b)
+      self.assertAllEqual(func_output.numpy(), [-2., -2.])
+
+      eager_output = model(a, b)
+      self.assertAllEqual(eager_output.numpy(), [-2., -2.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testSubGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(a)
+        tape.watch(b)
+        result = unified_math_ops.sub(a, b)
+      grads = tape.gradient(result, [a, b])
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1., 2.]))
+      b = TensorCastHelper(constant_op.constant([3., 4.]))
+
+      func_outputs = def_function.function(model)(a, b)
+      self.assertAllEqual(func_outputs[0].numpy(), [1.0, 1.0])
+      self.assertAllEqual(func_outputs[1].numpy(), [-1.0, -1.0])
+
+      eager_outputs = model(a, b)
+      self.assertAllEqual(eager_outputs[0].numpy(), [1.0, 1.0])
+      self.assertAllEqual(eager_outputs[1].numpy(), [-1.0, -1.0])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testMul(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      return unified_math_ops.mul(a, b)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1., 2.]))
+      b = TensorCastHelper(constant_op.constant([3., 4.]))
+
+      func_output = def_function.function(model)(a, b)
+      self.assertAllEqual(func_output.numpy(), [3., 8.])
+
+      eager_output = model(a, b)
+      self.assertAllEqual(eager_output.numpy(), [3., 8.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testMulGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(a)
+        tape.watch(b)
+        result = unified_math_ops.mul(a, b)
+      grads = tape.gradient(result, [a, b])
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1., 2.]))
+      b = TensorCastHelper(constant_op.constant([3., 4.]))
+
+      func_outputs = def_function.function(model)(a, b)
+      self.assertAllEqual(func_outputs[0].numpy(), [3., 4.])
+      self.assertAllEqual(func_outputs[1].numpy(), [1., 2.])
+
+      eager_outputs = model(a, b)
+      self.assertAllEqual(eager_outputs[0].numpy(), [3., 4.])
+      self.assertAllEqual(eager_outputs[1].numpy(), [1., 2.])
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testLog1p(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a):
+      return unified_math_ops.log1p(a)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1.]))
+
+      func_output = def_function.function(model)(a)
+      self.assertArrayNear(func_output.numpy(), [0.69314], 0.001)
+
+      eager_output = model(a)
+      self.assertArrayNear(eager_output.numpy(), [0.69314], 0.001)
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testLog1pGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(a)
+        result = unified_math_ops.log1p(a)
+      grads = tape.gradient(result, a)
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([1.]))
+
+      func_outputs = def_function.function(model)(a)
+      self.assertArrayNear(func_outputs.numpy(), [0.5], 0.001)
+
+      eager_outputs = model(a)
+      self.assertArrayNear(eager_outputs.numpy(), [0.5], 0.001)
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testDivNoNan(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      return unified_math_ops.div_no_nan(a, b)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([2.]))
+      b = TensorCastHelper(constant_op.constant([4.]))
+
+      func_output = def_function.function(model)(a, b)
+      self.assertArrayNear(func_output.numpy(), [0.5], 0.001)
+
+      eager_output = model(a, b)
+      self.assertArrayNear(eager_output.numpy(), [0.5], 0.001)
+
+  @parameterized.named_parameters([
+      ("Graph", False),
+      ("Mlir", True),
+  ])
+  def testDivNoNanGrad(self, use_mlir):
+    if use_mlir:
+      SetTracingImplementation("mlir")
+
+    def model(a, b):
+      with tape_lib.GradientTape() as tape:
+        tape.watch(a)
+        tape.watch(b)
+        result = unified_math_ops.div_no_nan(a, b)
+      grads = tape.gradient(result, [a, b])
+      return grads
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      a = TensorCastHelper(constant_op.constant([2.]))
+      b = TensorCastHelper(constant_op.constant([4.]))
+
+      func_outputs = def_function.function(model)(a, b)
+      self.assertArrayNear(func_outputs[0].numpy(), [0.25], 0.001)
+      self.assertArrayNear(func_outputs[1].numpy(), [-0.125], 0.001)
+
+      eager_outputs = model(a, b)
+      self.assertArrayNear(eager_outputs[0].numpy(), [0.25], 0.001)
+      self.assertArrayNear(eager_outputs[1].numpy(), [-0.125], 0.001)
+
+
+class UnifiedTapeBenchmark(test.Benchmark):
+
+  def _computeMnistMlpGrads(self, math_ops_lib, nn_ops_lib, backprop_lib, cast,
+                            num_iters, hidden_layers, hidden_size, batch_size):
+    batch_size = 1
+    image_size = 28 * 28
+    num_classes = 10
+
+    def model(x, hidden_weights, softmax_weight, labels):
+      with backprop_lib.GradientTape() as tape:
+        for weight in hidden_weights + [softmax_weight]:
+          tape.watch(weight)
+        for hidden_weight in hidden_weights:
+          x = math_ops_lib.mat_mul(x, hidden_weight)
+          x = nn_ops_lib.relu(x)
+        logits = math_ops_lib.mat_mul(x, softmax_weight)
+        loss = nn_ops_lib.sparse_softmax_cross_entropy_with_logits(
+            logits=logits, labels=labels)
+
+      grads = tape.gradient(loss, hidden_weights + [softmax_weight])
+      return grads
+
+    x = maybe_cast(array_ops.ones([batch_size, image_size]), cast)
+    hidden_weights = []
+    for i in range(hidden_layers):
+      hidden_weights.append(
+          maybe_cast(
+              random_ops.random_uniform(
+                  [hidden_size if i else image_size, hidden_size]), cast))
+    softmax_weight = maybe_cast(
+        random_ops.random_uniform([hidden_size, num_classes]), cast)
+    labels = maybe_cast(array_ops.zeros([batch_size], dtype=dtypes.int32), cast)
+
+    with context_lib.set_default(get_immediate_execution_context()):
+      # Warm up.
+      for _ in range(10):
+        model(x, hidden_weights, softmax_weight, labels)
+      runtimes = timeit.repeat(
+          lambda: model(x, hidden_weights, softmax_weight, labels),
+          repeat=num_iters,
+          number=10)
+    return min(runtimes) / 10
+
+  def benchmarkTwoHiddenLayerMnistEagerUnified(self):
+    num_iters = 100
+    duration = self._computeMnistMlpGrads(
+        unified_math_ops,
+        unified_nn_ops,
+        tape_lib,
+        True,
+        num_iters,
+        hidden_layers=2,
+        hidden_size=100,
+        batch_size=1)
+    self.report_benchmark(
+        name="TwoHiddenLayerMnistEagerUnified",
+        iters=num_iters,
+        wall_time=duration)
+
+  def benchmarkTwoHiddenLayerMnistEager(self):
+    num_iters = 100
+    duration = self._computeMnistMlpGrads(
+        math_ops,
+        nn_ops,
+        backprop,
+        False,
+        num_iters,
+        hidden_layers=2,
+        hidden_size=100,
+        batch_size=1)
+    self.report_benchmark(
+        name="TwoHiddenLayerMnistEager", iters=num_iters, wall_time=duration)
+
+  def benchmarkTenHiddenLayerMnistEagerUnified(self):
+    num_iters = 100
+    duration = self._computeMnistMlpGrads(
+        unified_math_ops,
+        unified_nn_ops,
+        tape_lib,
+        True,
+        num_iters,
+        hidden_layers=10,
+        hidden_size=100,
+        batch_size=1)
+    self.report_benchmark(
+        name="TenHiddenLayerMnistEagerUnified",
+        iters=num_iters,
+        wall_time=duration)
+
+  def benchmarkTenHiddenLayerMnistEager(self):
+    num_iters = 100
+    duration = self._computeMnistMlpGrads(
+        math_ops,
+        nn_ops,
+        backprop,
+        False,
+        num_iters,
+        hidden_layers=10,
+        hidden_size=100,
+        batch_size=1)
+    self.report_benchmark(
+        name="TenHiddenLayerMnistEager", iters=num_iters, wall_time=duration)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()