chore: import upstream snapshot with attribution

tensorflow/python/eager/backprop_util.py

@@ -0,0 +1,105 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Shared utilities related to backprop."""
+
+from tensorflow.core.config import flags
+from tensorflow.core.framework import types_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import handle_data_util
+from tensorflow.python.ops import math_ops
+
+
+def _DTypeFromTensor(tensor):
+  """Extract either `tensor.dtype` or the unanimous sub-type of a variant."""
+  dtype = tensor.dtype
+  if dtype.base_dtype == dtypes.variant:
+    # If we know statically that the data a variant points to is non-trainable
+    # then the variant itself is non-trainable.
+    if isinstance(tensor, ops.EagerTensor):
+      handle_data = tensor._handle_data  # pylint: disable=protected-access
+    else:
+      handle_data = handle_data_util.get_resource_handle_data(tensor)
+    if (handle_data is not None
+        and handle_data.is_set
+        and handle_data.shape_and_type):
+      first_type = handle_data.shape_and_type[0].dtype
+      # Some variants have statically unknown dtypes; we can't make inferences
+      # about trainability, so we conservatively assume they're trainable
+      # (which may waste memory passing zeros around, but will be correct).
+      if (first_type != types_pb2.DT_INVALID
+          and all(shape_and_type.dtype == first_type
+                  for shape_and_type in handle_data.shape_and_type)):
+        return first_type
+  return dtype
+
+
+def IsTrainable(tensor_or_dtype):
+  """Determines whether a tensor or dtype supports infinitesimal changes."""
+  if tensor_util.is_tf_type(tensor_or_dtype):
+    dtype = _DTypeFromTensor(tensor_or_dtype)
+  else:
+    dtype = tensor_or_dtype
+  dtype = dtypes.as_dtype(dtype)
+  trainable_dtypes = [dtypes.float16, dtypes.float32, dtypes.float64,
+                      dtypes.complex64, dtypes.complex128, dtypes.resource,
+                      dtypes.variant, dtypes.bfloat16]
+  if flags.config().enable_quantized_dtypes_training.value():
+    trainable_dtypes.extend([dtypes.qint8, dtypes.qint16, dtypes.qint32,
+                             dtypes.quint8, dtypes.quint16])
+  return dtype.base_dtype in trainable_dtypes
+
+
+def FlattenNestedIndexedSlices(grad):
+  assert isinstance(grad, indexed_slices.IndexedSlices)
+  if isinstance(grad.values, tensor_lib.Tensor):
+    return grad
+  else:
+    assert isinstance(grad.values, indexed_slices.IndexedSlices)
+    g = FlattenNestedIndexedSlices(grad.values)
+    return indexed_slices.IndexedSlices(
+        g.values, array_ops.gather(grad.indices, g.indices), g.dense_shape)
+
+
+def AggregateIndexedSlicesGradients(grads):
+  """Aggregates gradients containing `IndexedSlices`s."""
+  if len(grads) < 1:
+    return None
+  if len(grads) == 1:
+    return grads[0]
+  grads = [g for g in grads if g is not None]
+  # If any gradient is a `Tensor`, sum them up and return a dense tensor
+  # object.
+  if any(isinstance(g, tensor_lib.Tensor) for g in grads):
+    return math_ops.add_n(grads)
+
+  # The following `_as_indexed_slices_list` casts ids of IndexedSlices into
+  # int64. It is to make sure the inputs of `concat` all have same the data
+  # type.
+  grads = math_ops._as_indexed_slices_list(grads)  # pylint: disable=protected-access
+
+  grads = [FlattenNestedIndexedSlices(x) for x in grads]
+  # Form IndexedSlices out of the concatenated values and indices.
+  concat_grad = indexed_slices.IndexedSlices(
+      array_ops.concat([x.values for x in grads], axis=0),
+      array_ops.concat([x.indices for x in grads], axis=0),
+      grads[0].dense_shape)
+
+  return concat_grad
+

tensorflow/python/eager/benchmarks/BUILD

@@ -0,0 +1,27 @@
+load(
+    "//tensorflow/tools/test:performance.bzl",
+    "cuda_py_benchmark_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+cuda_py_benchmark_test(
+    name = "kpi_benchmark_test",
+    size = "medium",
+    srcs = ["kpi_benchmark_test.py"],
+    tags = [
+        "no_windows",  #  b/141617449
+        "notap",  # b/456542868
+        "optonly",
+    ],
+    deps = [
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/eager:benchmarks_test_base",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/profiler:trace",
+    ],
+)

tensorflow/python/eager/benchmarks/kpi_benchmark_test.py

@@ -0,0 +1,117 @@
+# 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.
+# ==============================================================================
+r"""KPI Benchmarks for low-level eager execution primitives.
+
+This is a suite of full end-to-end integration benchmakr for low-level eager
+execution APIs. Also tracks them as KPI Traceme.
+
+To run CPU benchmarks:
+  bazel run -c opt kpi_benchmarks_test -- --benchmark_filter=.
+
+To run GPU benchmarks:
+  bazel run --config=cuda -c opt --copt="-mavx" kpi_benchmarks_test -- \
+    --benchmark_filter=.
+
+To run a subset of benchmarks using --benchmarks flag.
+--benchmarks: the list of benchmarks to run. The specified value is interpreted
+as a regular expression and any benchmark whose name contains a partial match
+to the regular expression is executed.
+e.g. --benchmark_filter=".*matmul*." will run all matmul related benchmarks.
+
+"""
+import gc
+import time
+
+import tensorflow as tf
+
+from tensorflow.python.eager import benchmarks_test_base
+from tensorflow.python.eager import context
+from tensorflow.python.profiler import trace
+
+NUM_ITERATIONS = 30000
+
+
+def _run_benchmark(func, num_iters, execution_mode=None):
+  ctx = context.context()
+  with context.execution_mode(execution_mode):
+    # call func to warm up
+    func()
+    if execution_mode == context.ASYNC:
+      ctx.executor.wait()
+    start = time.time()
+    for _ in range(num_iters):
+      func()
+    if execution_mode == context.ASYNC:
+      ctx.executor.wait()
+    end = time.time()
+
+    return end - start
+
+
+class KpiBenchmarks(benchmarks_test_base.MicroBenchmarksBase):
+  """A Collection of KPI benchmarks."""
+
+  def _get_benchmark_name(self):
+    return self._get_name()
+
+  def _run(self, func, num_iters):
+    gc.disable()
+    gc.collect()
+    self.run_report(_run_benchmark, func, num_iters)
+    gc.enable()
+
+  def benchmark_tf_constant_2x2(self):
+    x = [[1., 2.], [3., 4.]]
+
+    def fn():
+      with trace.Trace("tf.constant-2x2"):
+        tf.constant(x)
+
+    self._run(fn, NUM_ITERATIONS)
+
+  def benchmark_tf_convert_to_tensor_2x2(self):
+    x = [[1., 2.], [3., 4.]]
+
+    def fn():
+      with trace.Trace("tf.convert_to_tensor-2x2"):
+        tf.convert_to_tensor(x)
+
+    self._run(fn, NUM_ITERATIONS)
+
+  def benchmark_tf_nn_relu_2x2(self):
+    x = tf.constant([[1., 2.], [3., 4.]])
+
+    def fn():
+      with trace.Trace("tf.nn.relu-2x2"):
+        tf.nn.relu(x)
+
+    self._run(fn, NUM_ITERATIONS)
+
+  def benchmark_tf_function_invocation_identity(self):
+    x = tf.constant([[1., 2.], [3., 4.]])
+
+    @tf.function
+    def identity(x):
+      return x
+
+    def fn():
+      with trace.Trace("tf.function-identity"):
+        identity(x)
+
+    self._run(fn, NUM_ITERATIONS)
+
+
+if __name__ == "__main__":
+  tf.test.main()

tensorflow/python/eager/benchmarks/resnet50/BUILD

@@ -0,0 +1,103 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "resnet50",
+    srcs = ["resnet50.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py_no_contrib",
+    ],
+)
+
+py_library(
+    name = "resnet50_test_util",
+    srcs = ["resnet50_test_util.py"],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py_no_contrib"],
+)
+
+py_library(
+    name = "resnet50_test_lib",
+    srcs = ["resnet50_test.py"],
+    strict_deps = True,
+    deps = [
+        ":resnet50",
+        ":resnet50_test_util",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/client:device_lib",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:record",
+        "//tensorflow/python/framework:test_lib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "resnet50_test",
+    size = "medium",
+    srcs = ["resnet50_test.py"],
+    shard_count = 4,
+    tags = [
+        "no_pip",
+        "no_windows",  # TODO(b/141617449): needs investigation
+        "optonly",
+        "oss_serial",
+        "v1only",
+    ],
+    deps = [
+        ":resnet50",
+        ":resnet50_test_util",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/client:device_lib",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:record",
+        "//tensorflow/python/framework:test_lib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "hvp_test",
+    size = "medium",
+    srcs = ["hvp_test.py"],
+    shard_count = 7,
+    tags = [
+        "no_pip",
+        "no_windows",  # TODO(b/141617449): needs investigation
+        "optonly",
+        "oss_serial",
+        "v1only",
+    ],
+    # Times out
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":resnet50",
+        ":resnet50_test_util",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/eager:forwardprop",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "resnet50_graph_test",
+    size = "medium",
+    srcs = ["resnet50_graph_test.py"],
+    shard_count = 4,
+    tags = [
+        "no_pip",
+        "no_windows",  # TODO(b/141617449): needs investigation
+        "optonly",
+        "oss_serial",
+    ],
+    deps = [
+        ":resnet50",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/python/eager/benchmarks/resnet50/README.md

@@ -0,0 +1,45 @@
+Image classification using the ResNet50 model described in
+[Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385).
+
+Contents:
+
+- `resnet50.py`: Model definition
+- `resnet50_test.py`: Sanity unittests and benchmarks for using the model with
+  eager execution enabled.
+- `resnet50_graph_test.py`: Sanity unittests and benchmarks when using the same
+  model code to construct a TensorFlow graph.
+
+# Benchmarks
+
+Using a synthetic data, run:
+
+```
+# Using eager execution
+python resnet50_test.py --benchmark_filter=.
+
+# Using graph execution
+python resnet50_graph_test.py --benchmark_filter=.
+```
+
+The above uses the model definition included with the TensorFlow pip
+package. To build (and run benchmarks) from source:
+
+```
+# Using eager execution
+bazel run -c opt --config=cuda :resnet50_test -- --benchmark_filter=.
+
+# Using graph execution
+bazel run -c opt --config=cuda :resnet50_graph_test -- --benchmark_filter=.
+```
+
+(Or remove the `--config=cuda` flag for running on CPU instead of GPU).
+
+On October 31, 2017, the benchmarks demonstrated comparable performance
+for eager and graph execution of this particular model when using
+a single NVIDIA Titan X (Pascal) GPU on a host with an
+Intel Xeon E5-1650 CPU @ 3.50GHz and a batch size of 32.
+
+| Benchmark name                           | batch size    | images/second |
+| ---------------------------------------  | ------------- | ------------- |
+| eager_train_gpu_batch_32_channels_first  |            32 |           171 |
+| graph_train_gpu_batch_32_channels_first  |            32 |           172 |

tensorflow/python/eager/benchmarks/resnet50/hvp_test.py

@@ -0,0 +1,177 @@
+# Copyright 2019 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 and benchmarks for Hessian-vector products with ResNet50."""
+
+import gc
+import time
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.python.eager import forwardprop
+from tensorflow.python.eager.benchmarks.resnet50 import resnet50
+from tensorflow.python.eager.benchmarks.resnet50 import resnet50_test_util
+
+
+def _forward_over_back_hvp(model, images, labels, vector):
+  with forwardprop.ForwardAccumulator(
+      model.trainable_variables, vector) as acc:
+    with tf.GradientTape() as grad_tape:
+      logits = model(images, training=True)
+      loss = tf.compat.v1.losses.softmax_cross_entropy(
+          logits=logits, onehot_labels=labels)
+    grads = grad_tape.gradient(loss, model.trainable_variables)
+  return acc.jvp(grads)
+
+
+def _back_over_forward_hvp(model, images, labels, vector):
+  with tf.GradientTape() as grad_tape:
+    grad_tape.watch(model.trainable_variables)
+    with forwardprop.ForwardAccumulator(
+        model.trainable_variables, vector) as acc:
+      logits = model(images, training=True)
+      loss = tf.compat.v1.losses.softmax_cross_entropy(
+          logits=logits, onehot_labels=labels)
+  return grad_tape.gradient(acc.jvp(loss), model.trainable_variables)
+
+
+def _tf_gradients_forward_over_back_hvp(model, images, labels, vector):
+  with tf.GradientTape() as grad_tape:
+    logits = model(images, training=True)
+    loss = tf.compat.v1.losses.softmax_cross_entropy(
+        logits=logits, onehot_labels=labels)
+  variables = model.trainable_variables
+  grads = grad_tape.gradient(loss, variables)
+  helpers = tf.nest.map_structure(tf.ones_like, grads)
+  transposing = tf.gradients(grads, variables, helpers)
+  return tf.gradients(transposing, helpers, vector)
+
+
+def _back_over_back_hvp(model, images, labels, vector):
+  with tf.GradientTape() as outer_tape:
+    with tf.GradientTape() as inner_tape:
+      logits = model(images, training=True)
+      loss = tf.compat.v1.losses.softmax_cross_entropy(
+          logits=logits, onehot_labels=labels)
+    grads = inner_tape.gradient(loss, model.trainable_variables)
+  return outer_tape.gradient(
+      grads, model.trainable_variables, output_gradients=vector)
+
+
+class HVPTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ("forward_over_back_eager", _forward_over_back_hvp),
+      ("forward_over_back_function", tf.function(_forward_over_back_hvp)),
+      ("tf_gradients", tf.function(_tf_gradients_forward_over_back_hvp)),
+      ("back_over_back_eager", _back_over_back_hvp),
+      ("back_over_back_function", tf.function(_back_over_back_hvp)),
+      ("back_over_forward_eager", _back_over_forward_hvp),
+      ("back_over_forward_function", tf.function(_back_over_forward_hvp)))
+  def test_hvp_shapes(self, hvp_function):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format)
+    with tf.device(device):
+      images, labels = resnet50_test_util.random_batch(2, data_format)
+      images = tf.constant(images)
+      labels = tf.constant(labels)
+      model.build(images.shape)
+      vector = [tf.ones_like(v) for v in model.trainable_variables]
+
+      # Note that numerical differences build up to quite large differences here
+      # in the final hvp. tensorflow/python/eager:forwardprop_test has a
+      # smaller-scale test that the computations are close on a much smaller but
+      # otherwise similar model.
+      hvp = hvp_function(model, images, labels, vector)
+      for hvp_component, variable in zip(hvp, model.trainable_variables):
+        self.assertEqual(hvp_component.shape, variable.shape)
+        self.assertEqual(hvp_component.dtype, variable.dtype)
+
+
+class HVPBenchmarks(tf.test.Benchmark):
+
+  def _force_device_sync(self):
+    # If this function is called in the context of a non-CPU device
+    # (e.g., inside a 'with tf.device("/gpu:0")' block)
+    # then this will force a copy from CPU->NON_CPU_DEVICE->CPU,
+    # which forces a sync. This is a roundabout way, yes.
+    tf.constant(1.).cpu()
+
+  def _hvp_benchmark(self, hvp_fn, label, batch_sizes,
+                     num_iters=30, num_burn=5):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format)
+    for batch_size in batch_sizes:
+      with tf.device(device):
+        images, labels = resnet50_test_util.random_batch(
+            batch_size, data_format)
+        images = tf.constant(images)
+        labels = tf.constant(labels)
+        model.build(images.shape)
+        vector = [tf.ones_like(v) for v in model.trainable_variables]
+        for _ in range(num_burn):
+          results = hvp_fn(model, images, labels, vector)
+          for result in results:
+            result.cpu()
+        self._force_device_sync()
+        gc.collect()
+        start = time.time()
+        for _ in range(num_iters):
+          results = hvp_fn(model, images, labels, vector)
+          for result in results:
+            result.cpu()
+        self._force_device_sync()
+        resnet50_test_util.report(
+            self, label, start, num_iters, device, batch_size, data_format)
+
+  def benchmark_forward_over_backward_hvp_eager(self):
+    self._hvp_benchmark(_forward_over_back_hvp,
+                        "forward_over_backward_hvp_eager",
+                        batch_sizes=[8])
+
+  def benchmark_forward_over_backward_hvp_function(self):
+    self._hvp_benchmark(tf.function(_forward_over_back_hvp),
+                        "forward_over_backward_hvp_function",
+                        batch_sizes=[8])
+
+  def benchmark_tf_gradients_forward_over_backward_hvp_function(self):
+    self._hvp_benchmark(tf.function(_tf_gradients_forward_over_back_hvp),
+                        "tf_gradients_forward_over_backward_hvp_function",
+                        batch_sizes=[8])
+
+  def benchmark_backward_over_backward_hvp_eager(self):
+    self._hvp_benchmark(_back_over_back_hvp,
+                        "backward_over_backward_hvp_eager",
+                        batch_sizes=[8])
+
+  def benchmark_backward_over_backward_hvp_function(self):
+    self._hvp_benchmark(tf.function(_back_over_back_hvp),
+                        "backward_over_backward_hvp_function",
+                        batch_sizes=[8])
+
+  def benchmark_backward_over_forward_hvp_eager(self):
+    self._hvp_benchmark(_back_over_forward_hvp,
+                        "backward_over_forward_hvp_eager",
+                        batch_sizes=[8])
+
+  def benchmark_backward_over_forward_hvp_function(self):
+    self._hvp_benchmark(tf.function(_back_over_forward_hvp),
+                        "backward_over_forward_hvp_function",
+                        batch_sizes=[8])
+
+
+if __name__ == "__main__":
+  tf.compat.v1.enable_v2_behavior()
+  tf.test.main()

tensorflow/python/eager/benchmarks/resnet50/resnet50.py

@@ -0,0 +1,365 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""ResNet50 model definition compatible with TensorFlow's eager execution.
+
+Reference [Deep Residual Learning for Image
+Recognition](https://arxiv.org/abs/1512.03385)
+
+Adapted from tf.keras.applications.ResNet50. A notable difference is that the
+model here outputs logits while the Keras model outputs probability.
+"""
+import functools
+
+import tensorflow as tf
+
+layers = tf.keras.layers
+
+
+class _IdentityBlock(tf.keras.Model):
+  """_IdentityBlock is the block that has no conv layer at shortcut.
+
+  Args:
+    kernel_size: the kernel size of middle conv layer at main path
+    filters: list of integers, the filters of 3 conv layer at main path
+    stage: integer, current stage label, used for generating layer names
+    block: 'a','b'..., current block label, used for generating layer names
+    data_format: data_format for the input ('channels_first' or
+      'channels_last').
+  """
+
+  def __init__(self, kernel_size, filters, stage, block, data_format):
+    super(_IdentityBlock, self).__init__(name='')
+    filters1, filters2, filters3 = filters
+
+    conv_name_base = 'res' + str(stage) + block + '_branch'
+    bn_name_base = 'bn' + str(stage) + block + '_branch'
+    bn_axis = 1 if data_format == 'channels_first' else 3
+
+    self.conv2a = layers.Conv2D(
+        filters1, (1, 1), name=conv_name_base + '2a', data_format=data_format)
+    self.bn2a = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '2a')
+
+    self.conv2b = layers.Conv2D(
+        filters2,
+        kernel_size,
+        padding='same',
+        data_format=data_format,
+        name=conv_name_base + '2b')
+    self.bn2b = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '2b')
+
+    self.conv2c = layers.Conv2D(
+        filters3, (1, 1), name=conv_name_base + '2c', data_format=data_format)
+    self.bn2c = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '2c')
+
+  def call(self, input_tensor, training=False):
+    x = self.conv2a(input_tensor)
+    x = self.bn2a(x, training=training)
+    x = tf.nn.relu(x)
+
+    x = self.conv2b(x)
+    x = self.bn2b(x, training=training)
+    x = tf.nn.relu(x)
+
+    x = self.conv2c(x)
+    x = self.bn2c(x, training=training)
+
+    x += input_tensor
+    return tf.nn.relu(x)
+
+
+class _ConvBlock(tf.keras.Model):
+  """_ConvBlock is the block that has a conv layer at shortcut.
+
+  Args:
+      kernel_size: the kernel size of middle conv layer at main path
+      filters: list of integers, the filters of 3 conv layer at main path
+      stage: integer, current stage label, used for generating layer names
+      block: 'a','b'..., current block label, used for generating layer names
+      data_format: data_format for the input ('channels_first' or
+        'channels_last').
+      strides: strides for the convolution. Note that from stage 3, the first
+       conv layer at main path is with strides=(2,2), and the shortcut should
+       have strides=(2,2) as well.
+  """
+
+  def __init__(self,
+               kernel_size,
+               filters,
+               stage,
+               block,
+               data_format,
+               strides=(2, 2)):
+    super(_ConvBlock, self).__init__(name='')
+    filters1, filters2, filters3 = filters
+
+    conv_name_base = 'res' + str(stage) + block + '_branch'
+    bn_name_base = 'bn' + str(stage) + block + '_branch'
+    bn_axis = 1 if data_format == 'channels_first' else 3
+
+    self.conv2a = layers.Conv2D(
+        filters1, (1, 1),
+        strides=strides,
+        name=conv_name_base + '2a',
+        data_format=data_format)
+    self.bn2a = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '2a')
+
+    self.conv2b = layers.Conv2D(
+        filters2,
+        kernel_size,
+        padding='same',
+        name=conv_name_base + '2b',
+        data_format=data_format)
+    self.bn2b = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '2b')
+
+    self.conv2c = layers.Conv2D(
+        filters3, (1, 1), name=conv_name_base + '2c', data_format=data_format)
+    self.bn2c = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '2c')
+
+    self.conv_shortcut = layers.Conv2D(
+        filters3, (1, 1),
+        strides=strides,
+        name=conv_name_base + '1',
+        data_format=data_format)
+    self.bn_shortcut = layers.BatchNormalization(
+        axis=bn_axis, name=bn_name_base + '1')
+
+  def call(self, input_tensor, training=False):
+    x = self.conv2a(input_tensor)
+    x = self.bn2a(x, training=training)
+    x = tf.nn.relu(x)
+
+    x = self.conv2b(x)
+    x = self.bn2b(x, training=training)
+    x = tf.nn.relu(x)
+
+    x = self.conv2c(x)
+    x = self.bn2c(x, training=training)
+
+    shortcut = self.conv_shortcut(input_tensor)
+    shortcut = self.bn_shortcut(shortcut, training=training)
+
+    x += shortcut
+    return tf.nn.relu(x)
+
+
+# pylint: disable=not-callable
+class ResNet50(tf.keras.Model):
+  """Instantiates the ResNet50 architecture.
+
+  Args:
+    data_format: format for the image. Either 'channels_first' or
+      'channels_last'.  'channels_first' is typically faster on GPUs while
+      'channels_last' is typically faster on CPUs. See
+      https://www.tensorflow.org/performance/performance_guide#data_formats
+    name: Prefix applied to names of variables created in the model.
+    trainable: Is the model trainable? If true, performs backward
+        and optimization after call() method.
+    include_top: whether to include the fully-connected layer at the top of the
+      network.
+    pooling: Optional pooling mode for feature extraction when `include_top`
+      is `False`.
+      - `None` means that the output of the model will be the 4D tensor
+          output of the last convolutional layer.
+      - `avg` means that global average pooling will be applied to the output of
+          the last convolutional layer, and thus the output of the model will be
+          a 2D tensor.
+      - `max` means that global max pooling will be applied.
+    block3_strides: whether to add a stride of 2 to block3 to make it compatible
+      with tf.slim ResNet implementation.
+    average_pooling: whether to do average pooling of block4 features before
+      global pooling.
+    classes: optional number of classes to classify images into, only to be
+      specified if `include_top` is True.
+
+  Raises:
+      ValueError: in case of invalid argument for data_format.
+  """
+
+  def __init__(self,
+               data_format,
+               name='',
+               trainable=True,
+               include_top=True,
+               pooling=None,
+               block3_strides=False,
+               average_pooling=True,
+               classes=1000):
+    super(ResNet50, self).__init__(name=name)
+
+    valid_channel_values = ('channels_first', 'channels_last')
+    if data_format not in valid_channel_values:
+      raise ValueError('Unknown data_format: %s. Valid values: %s' %
+                       (data_format, valid_channel_values))
+    self.include_top = include_top
+    self.block3_strides = block3_strides
+    self.average_pooling = average_pooling
+    self.pooling = pooling
+
+    def conv_block(filters, stage, block, strides=(2, 2)):
+      return _ConvBlock(
+          3,
+          filters,
+          stage=stage,
+          block=block,
+          data_format=data_format,
+          strides=strides)
+
+    def id_block(filters, stage, block):
+      return _IdentityBlock(
+          3, filters, stage=stage, block=block, data_format=data_format)
+
+    self.conv1 = layers.Conv2D(
+        64, (7, 7),
+        strides=(2, 2),
+        data_format=data_format,
+        padding='same',
+        name='conv1')
+    bn_axis = 1 if data_format == 'channels_first' else 3
+    self.bn_conv1 = layers.BatchNormalization(axis=bn_axis, name='bn_conv1')
+    self.max_pool = layers.MaxPooling2D((3, 3),
+                                        strides=(2, 2),
+                                        data_format=data_format)
+
+    self.l2a = conv_block([64, 64, 256], stage=2, block='a', strides=(1, 1))
+    self.l2b = id_block([64, 64, 256], stage=2, block='b')
+    self.l2c = id_block([64, 64, 256], stage=2, block='c')
+
+    self.l3a = conv_block([128, 128, 512], stage=3, block='a')
+    self.l3b = id_block([128, 128, 512], stage=3, block='b')
+    self.l3c = id_block([128, 128, 512], stage=3, block='c')
+    self.l3d = id_block([128, 128, 512], stage=3, block='d')
+
+    self.l4a = conv_block([256, 256, 1024], stage=4, block='a')
+    self.l4b = id_block([256, 256, 1024], stage=4, block='b')
+    self.l4c = id_block([256, 256, 1024], stage=4, block='c')
+    self.l4d = id_block([256, 256, 1024], stage=4, block='d')
+    self.l4e = id_block([256, 256, 1024], stage=4, block='e')
+    self.l4f = id_block([256, 256, 1024], stage=4, block='f')
+
+    # Striding layer that can be used on top of block3 to produce feature maps
+    # with the same resolution as the TF-Slim implementation.
+    if self.block3_strides:
+      self.subsampling_layer = layers.MaxPooling2D((1, 1),
+                                                   strides=(2, 2),
+                                                   data_format=data_format)
+      self.l5a = conv_block([512, 512, 2048],
+                            stage=5,
+                            block='a',
+                            strides=(1, 1))
+    else:
+      self.l5a = conv_block([512, 512, 2048], stage=5, block='a')
+    self.l5b = id_block([512, 512, 2048], stage=5, block='b')
+    self.l5c = id_block([512, 512, 2048], stage=5, block='c')
+
+    self.avg_pool = layers.AveragePooling2D((7, 7),
+                                            strides=(7, 7),
+                                            data_format=data_format)
+
+    if self.include_top:
+      self.flatten = layers.Flatten()
+      self.fc1000 = layers.Dense(classes, name='fc1000')
+    else:
+      reduction_indices = [1, 2] if data_format == 'channels_last' else [2, 3]
+      reduction_indices = tf.constant(reduction_indices)
+      if pooling == 'avg':
+        self.global_pooling = functools.partial(
+            tf.reduce_mean,
+            axis=reduction_indices,
+            keepdims=False)
+      elif pooling == 'max':
+        self.global_pooling = functools.partial(
+            tf.reduce_max, reduction_indices=reduction_indices, keep_dims=False)
+      else:
+        self.global_pooling = None
+
+  def call(self, inputs, training=True, intermediates_dict=None):
+    """Call the ResNet50 model.
+
+    Args:
+      inputs: Images to compute features for.
+      training: Whether model is in training phase.
+      intermediates_dict: `None` or dictionary. If not None, accumulate feature
+        maps from intermediate blocks into the dictionary.
+        ""
+
+    Returns:
+      Tensor with featuremap.
+    """
+
+    x = self.conv1(inputs)
+    x = self.bn_conv1(x, training=training)
+    x = tf.nn.relu(x)
+    if intermediates_dict is not None:
+      intermediates_dict['block0'] = x
+
+    x = self.max_pool(x)
+    if intermediates_dict is not None:
+      intermediates_dict['block0mp'] = x
+
+    # Block 1 (equivalent to "conv2" in Resnet paper).
+    x = self.l2a(x, training=training)
+    x = self.l2b(x, training=training)
+    x = self.l2c(x, training=training)
+    if intermediates_dict is not None:
+      intermediates_dict['block1'] = x
+
+    # Block 2 (equivalent to "conv3" in Resnet paper).
+    x = self.l3a(x, training=training)
+    x = self.l3b(x, training=training)
+    x = self.l3c(x, training=training)
+    x = self.l3d(x, training=training)
+    if intermediates_dict is not None:
+      intermediates_dict['block2'] = x
+
+    # Block 3 (equivalent to "conv4" in Resnet paper).
+    x = self.l4a(x, training=training)
+    x = self.l4b(x, training=training)
+    x = self.l4c(x, training=training)
+    x = self.l4d(x, training=training)
+    x = self.l4e(x, training=training)
+    x = self.l4f(x, training=training)
+
+    if self.block3_strides:
+      x = self.subsampling_layer(x)
+      if intermediates_dict is not None:
+        intermediates_dict['block3'] = x
+    else:
+      if intermediates_dict is not None:
+        intermediates_dict['block3'] = x
+
+    x = self.l5a(x, training=training)
+    x = self.l5b(x, training=training)
+    x = self.l5c(x, training=training)
+
+    if self.average_pooling:
+      x = self.avg_pool(x)
+      if intermediates_dict is not None:
+        intermediates_dict['block4'] = x
+    else:
+      if intermediates_dict is not None:
+        intermediates_dict['block4'] = x
+
+    if self.include_top:
+      return self.fc1000(self.flatten(x))
+    elif self.global_pooling:
+      return self.global_pooling(x)
+    else:
+      return x

tensorflow/python/eager/benchmarks/resnet50/resnet50_graph_test.py

@@ -0,0 +1,126 @@
+# Copyright 2017 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 and benchmarks for ResNet50 under graph execution."""
+
+import time
+
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from tensorflow.python.eager.benchmarks.resnet50 import resnet50
+
+
+def data_format():
+  return 'channels_first' if tf.test.is_gpu_available() else 'channels_last'
+
+
+def image_shape(batch_size):
+  if data_format() == 'channels_first':
+    return [batch_size, 3, 224, 224]
+  return [batch_size, 224, 224, 3]
+
+
+def random_batch(batch_size):
+  images = np.random.rand(*image_shape(batch_size)).astype(np.float32)
+  num_classes = 1000
+  labels = np.random.randint(
+      low=0, high=num_classes, size=[batch_size]).astype(np.int32)
+  one_hot = np.zeros((batch_size, num_classes)).astype(np.float32)
+  one_hot[np.arange(batch_size), labels] = 1.
+  return images, one_hot
+
+
+class ResNet50GraphTest(tf.test.TestCase):
+
+  def testApply(self):
+    # Use small batches for tests because the OSS version runs
+    # in constrained GPU environment with 1-2GB of memory.
+    batch_size = 8
+    with tf.Graph().as_default():
+      images = tf.placeholder(tf.float32, image_shape(None))
+      model = resnet50.ResNet50(data_format())
+      predictions = model(images, training=False)
+
+      init = tf.global_variables_initializer()
+
+      with tf.Session() as sess:
+        sess.run(init)
+        np_images, _ = random_batch(batch_size)
+        out = sess.run(predictions, feed_dict={images: np_images})
+        self.assertAllEqual([batch_size, 1000], out.shape)
+
+
+class ResNet50Benchmarks(tf.test.Benchmark):
+
+  def _report(self, label, start, num_iters, batch_size):
+    avg_time = (time.time() - start) / num_iters
+    dev = 'gpu' if tf.test.is_gpu_available() else 'cpu'
+    name = 'graph_%s_%s_batch_%d_%s' % (label, dev, batch_size, data_format())
+    extras = {'examples_per_sec': batch_size / avg_time}
+    self.report_benchmark(
+        iters=num_iters, wall_time=avg_time, name=name, extras=extras)
+
+  def benchmark_graph_apply(self):
+    with tf.Graph().as_default():
+      images = tf.placeholder(tf.float32, image_shape(None))
+      model = resnet50.ResNet50(data_format())
+      predictions = model(images, training=False)
+
+      init = tf.global_variables_initializer()
+
+      batch_size = 64
+      with tf.Session() as sess:
+        sess.run(init)
+        np_images, _ = random_batch(batch_size)
+        num_burn, num_iters = (3, 30)
+        for _ in range(num_burn):
+          sess.run(predictions, feed_dict={images: np_images})
+        start = time.time()
+        for _ in range(num_iters):
+          # Comparison with the eager execution benchmark in resnet50_test.py
+          # isn't entirely fair as the time here includes the cost of copying
+          # the feeds from CPU memory to GPU.
+          sess.run(predictions, feed_dict={images: np_images})
+        self._report('apply', start, num_iters, batch_size)
+
+  def benchmark_graph_train(self):
+    for batch_size in [16, 32, 64]:
+      with tf.Graph().as_default():
+        np_images, np_labels = random_batch(batch_size)
+        dataset = tf.data.Dataset.from_tensors((np_images, np_labels)).repeat()
+        images, labels = tf.data.make_one_shot_iterator(
+            dataset).get_next()
+
+        model = resnet50.ResNet50(data_format())
+        logits = model(images, training=True)
+        loss = tf.compat.v1.losses.softmax_cross_entropy(
+            logits=logits, onehot_labels=labels)
+        optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0)
+        train_op = optimizer.minimize(loss)
+
+        init = tf.global_variables_initializer()
+        with tf.Session() as sess:
+          sess.run(init)
+          (num_burn, num_iters) = (5, 10)
+          for _ in range(num_burn):
+            sess.run(train_op)
+          start = time.time()
+          for _ in range(num_iters):
+            sess.run(train_op)
+          self._report('train', start, num_iters, batch_size)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/eager/benchmarks/resnet50/resnet50_test.py

@@ -0,0 +1,427 @@
+# Copyright 2017 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 and benchmarks for the ResNet50 model, executed eagerly."""
+
+import gc
+import os
+import tempfile
+import time
+
+import tensorflow as tf
+
+from tensorflow.python.client import device_lib
+from tensorflow.python.eager import context
+from tensorflow.python.eager import record
+from tensorflow.python.eager.benchmarks.resnet50 import resnet50
+from tensorflow.python.eager.benchmarks.resnet50 import resnet50_test_util
+from tensorflow.python.framework import test_util
+
+
+def compute_gradients(model, images, labels, num_replicas=1):
+  with tf.GradientTape() as grad_tape:
+    logits = model(images, training=True)
+    loss = tf.compat.v1.losses.softmax_cross_entropy(
+        logits=logits, onehot_labels=labels)
+    tf.compat.v2.summary.write('loss', loss)
+    if num_replicas != 1:
+      loss /= num_replicas
+
+  # TODO(b/110991947): We can mistakenly trace the gradient call in
+  # multi-threaded environment. Explicitly disable recording until
+  # this is fixed.
+  with record.stop_recording():
+    grads = grad_tape.gradient(loss, model.variables)
+  return grads
+
+
+def apply_gradients(model, optimizer, gradients):
+  optimizer.apply_gradients(zip(gradients, model.variables))
+
+
+def _events_from_file(filepath):
+  """Returns all events in a single event file.
+
+  Args:
+    filepath: Path to the event file.
+
+  Returns:
+    A list of all tf.compat.v1.Event protos in the event file.
+  """
+  records = list(tf.compat.v1.python_io.tf_record_iterator(filepath))
+  result = []
+  for r in records:
+    event = tf.compat.v1.Event()
+    event.ParseFromString(r)
+    result.append(event)
+  return result
+
+
+def events_from_logdir(logdir):
+  """Returns all events in the single eventfile in logdir.
+
+  Args:
+    logdir: The directory in which the single event file is sought.
+
+  Returns:
+    A list of all tf.compat.v1.Event protos from the single event file.
+
+  Raises:
+    AssertionError: If logdir does not contain exactly one file.
+  """
+  assert tf.io.gfile.exists(logdir)
+  files = tf.io.gfile.listdir(logdir)
+  assert len(files) == 1, 'Found not exactly one file in logdir: %s' % files
+  return _events_from_file(os.path.join(logdir, files[0]))
+
+
+class ResNet50Test(tf.test.TestCase):
+
+  def _apply(self, defun=False, execution_mode=None):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format)
+    if defun:
+      model.call = tf.function(model.call)
+    with tf.device(device), context.execution_mode(execution_mode):
+      images, _ = resnet50_test_util.random_batch(2, data_format)
+      output = model(images, training=False)
+      context.async_wait()
+    self.assertEqual((2, 1000), output.shape)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_apply(self):
+    self._apply(defun=False)
+
+  def test_apply_async(self):
+    self._apply(defun=False, execution_mode=context.ASYNC)
+
+  def test_apply_with_defun(self):
+    self._apply(defun=True)
+
+  def test_apply_with_defun_async(self):
+    self._apply(defun=True, execution_mode=context.ASYNC)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_apply_no_top(self):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format, include_top=False)
+    with tf.device(device):
+      images, _ = resnet50_test_util.random_batch(2, data_format)
+      output = model(images, training=False)
+    output_shape = ((2, 2048, 1, 1)
+                    if data_format == 'channels_first' else (2, 1, 1, 2048))
+    self.assertEqual(output_shape, output.shape)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_apply_with_pooling(self):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format, include_top=False, pooling='avg')
+    with tf.device(device):
+      images, _ = resnet50_test_util.random_batch(2, data_format)
+      output = model(images, training=False)
+    self.assertEqual((2, 2048), output.shape)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_apply_no_average_pooling(self):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(
+        data_format, average_pooling=False, include_top=False)
+    with tf.device(device):
+      images, _ = resnet50_test_util.random_batch(2, data_format)
+      output = model(images, training=False)
+    output_shape = ((2, 2048, 7, 7) if data_format == 'channels_first' else
+                    (2, 7, 7, 2048))
+    self.assertEqual(output_shape, output.shape)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_apply_block3_strides(self):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(
+        data_format, block3_strides=True, include_top=False)
+    with tf.device(device):
+      images, _ = resnet50_test_util.random_batch(2, data_format)
+      output = model(images, training=False)
+    output_shape = ((2, 2048, 1, 1) if data_format == 'channels_first' else
+                    (2, 1, 1, 2048))
+    self.assertEqual(output_shape, output.shape)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_apply_retrieve_intermediates(self):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(
+        data_format, block3_strides=True, include_top=False)
+    intermediates_dict = {}
+    with tf.device(device):
+      images, _ = resnet50_test_util.random_batch(2, data_format)
+      output = model(images, training=False,
+                     intermediates_dict=intermediates_dict)
+    output_shape = ((2, 2048, 1, 1) if data_format == 'channels_first' else
+                    (2, 1, 1, 2048))
+    self.assertEqual(output_shape, output.shape)
+
+    if data_format == 'channels_first':
+      block_shapes = {
+          'block0': (2, 64, 112, 112),
+          'block0mp': (2, 64, 55, 55),
+          'block1': (2, 256, 55, 55),
+          'block2': (2, 512, 28, 28),
+          'block3': (2, 1024, 7, 7),
+          'block4': (2, 2048, 1, 1),
+      }
+    else:
+      block_shapes = {
+          'block0': (2, 112, 112, 64),
+          'block0mp': (2, 55, 55, 64),
+          'block1': (2, 55, 55, 256),
+          'block2': (2, 28, 28, 512),
+          'block3': (2, 7, 7, 1024),
+          'block4': (2, 1, 1, 2048),
+      }
+    for (block_name, block) in intermediates_dict.items():
+      self.assertEqual(block_shapes[block_name], block.shape)
+
+  def _test_train(self, execution_mode=None):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format)
+    tf.compat.v2.summary.experimental.set_step(
+        tf.compat.v1.train.get_or_create_global_step())
+    logdir = tempfile.mkdtemp()
+    with tf.compat.v2.summary.create_file_writer(
+        logdir, max_queue=0,
+        name='t0').as_default(), tf.compat.v2.summary.record_if(True):
+      with tf.device(device), context.execution_mode(execution_mode):
+        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.1)
+        images, labels = resnet50_test_util.random_batch(2, data_format)
+        apply_gradients(model, optimizer,
+                        compute_gradients(model, images, labels))
+        self.assertEqual(320, len(model.variables))
+        context.async_wait()
+    events = events_from_logdir(logdir)
+    self.assertEqual(len(events), 2)
+    self.assertEqual(events[1].summary.value[0].tag, 'loss')
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_train(self):
+    self._test_train()
+
+  @test_util.disable_tfrt('TFE_ContextGetExecutorForThread missing b/156188669')
+  def test_train_async(self):
+    self._test_train(execution_mode=context.ASYNC)
+
+  @test_util.disable_tfrt('Flaky test. b/157103729')
+  def test_no_garbage(self):
+    device, data_format = resnet50_test_util.device_and_data_format()
+    model = resnet50.ResNet50(data_format)
+    optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.1)
+    with tf.device(device):
+      images, labels = resnet50_test_util.random_batch(2, data_format)
+      gc.disable()
+      # Warm up. Note that this first run does create significant amounts of
+      # garbage to be collected. The hope is that this is a build-only effect,
+      # and a subsequent training loop will create nothing which needs to be
+      # collected.
+      apply_gradients(model, optimizer,
+                      compute_gradients(model, images, labels))
+      gc.collect()
+      previous_gc_debug_flags = gc.get_debug()
+      gc.set_debug(gc.DEBUG_SAVEALL)
+      for _ in range(2):
+        # Run twice to ensure that garbage that is created on the first
+        # iteration is no longer accessible.
+        apply_gradients(model, optimizer,
+                        compute_gradients(model, images, labels))
+      gc.collect()
+      # There should be no garbage requiring collection.
+      self.assertEqual(0, len(gc.garbage))
+      gc.set_debug(previous_gc_debug_flags)
+      gc.enable()
+
+
+class MockIterator(object):
+
+  def __init__(self, tensors):
+    self._tensors = [tf.identity(x) for x in tensors]
+
+  def next(self):
+    return self._tensors
+
+
+class ResNet50Benchmarks(tf.test.Benchmark):
+
+  def _report(self, label, start, num_iters, device, batch_size, data_format,
+              num_replicas=1):
+    resnet50_test_util.report(self, label, start, num_iters, device, batch_size,
+                              data_format, num_replicas)
+
+  def _train_batch_sizes(self):
+    """Choose batch sizes based on GPU capability."""
+    for device in device_lib.list_local_devices():
+      # TODO(b/141475121): We need some way to check which batch sizes would
+      # work using a public API.
+      if tf.DeviceSpec.from_string(device.name).device_type == 'GPU':
+        # Avoid OOM errors with larger batch sizes, which seem to cause errors
+        # later on even if caught.
+        #
+        # TODO(allenl): Base this on device memory; memory limit information
+        # during the test seems to exclude the amount TensorFlow has allocated,
+        # which isn't useful.
+        if 'K20' in device.physical_device_desc:
+          return (16,)
+        # Quardro P1000.
+        if 'P1000' in device.physical_device_desc:
+          return (16,)
+        if 'P100' in device.physical_device_desc:
+          return (16, 32, 64)
+
+      if tf.DeviceSpec.from_string(device.name).device_type == 'TPU':
+        return (32,)
+    return (16, 32)
+
+  def _force_device_sync(self):
+    # If this function is called in the context of a non-CPU device
+    # (e.g., inside a 'with tf.device("/gpu:0")' block)
+    # then this will force a copy from CPU->NON_CPU_DEVICE->CPU,
+    # which forces a sync. This is a roundabout way, yes.
+    tf.constant(1.).cpu()
+
+  def _benchmark_eager_apply(self, label, device_and_format, defun=False,
+                             execution_mode=None):
+    with context.execution_mode(execution_mode):
+      device, data_format = device_and_format
+      model = resnet50.ResNet50(data_format)
+      if defun:
+        model.call = tf.function(model.call)
+      batch_size = 64
+      num_burn = 5
+      num_iters = 30
+      with tf.device(device):
+        images, _ = resnet50_test_util.random_batch(batch_size, data_format)
+        for _ in range(num_burn):
+          model(images, training=False).cpu()
+        if execution_mode:
+          context.async_wait()
+        gc.collect()
+        start = time.time()
+        for _ in range(num_iters):
+          model(images, training=False).cpu()
+        if execution_mode:
+          context.async_wait()
+        self._report(label, start, num_iters, device, batch_size, data_format)
+
+  def benchmark_eager_apply_sync(self):
+    self._benchmark_eager_apply(
+        'eager_apply', resnet50_test_util.device_and_data_format(),
+        defun=False)
+
+  def benchmark_eager_apply_async(self):
+    self._benchmark_eager_apply(
+        'eager_apply_async',
+        resnet50_test_util.device_and_data_format(),
+        defun=False,
+        execution_mode=context.ASYNC)
+
+  def benchmark_eager_apply_with_defun(self):
+    self._benchmark_eager_apply(
+        'eager_apply_with_defun',
+        resnet50_test_util.device_and_data_format(), defun=True)
+
+  def _benchmark_eager_train(self,
+                             label,
+                             make_iterator,
+                             device_and_format,
+                             defun=False,
+                             execution_mode=None):
+    with context.execution_mode(execution_mode):
+      device, data_format = device_and_format
+      for batch_size in self._train_batch_sizes():
+        (images, labels) = resnet50_test_util.random_batch(
+            batch_size, data_format)
+        model = resnet50.ResNet50(data_format)
+        # TODO(b/161911585): tf_to_corert MLIR lowering pipeline should handle
+        # case when momentum is not set.
+        optimizer = tf.keras.optimizers.SGD(0.1, 0.1)
+        apply_grads = apply_gradients
+        if defun:
+          model.call = tf.function(model.call)
+          apply_grads = tf.function(apply_gradients)
+
+        num_burn = 3
+        num_iters = 10
+        with tf.device(device):
+          iterator = make_iterator((images, labels))
+          for _ in range(num_burn):
+            (images, labels) = iterator.next()
+            apply_grads(model, optimizer,
+                        compute_gradients(model, images, labels))
+          if execution_mode:
+            context.async_wait()
+          self._force_device_sync()
+          gc.collect()
+
+          start = time.time()
+          for _ in range(num_iters):
+            (images, labels) = iterator.next()
+            apply_grads(model, optimizer,
+                        compute_gradients(model, images, labels))
+          if execution_mode:
+            context.async_wait()
+          self._force_device_sync()
+          self._report(label, start, num_iters, device, batch_size, data_format)
+
+  def benchmark_eager_train_sync(self):
+    self._benchmark_eager_train(
+        'eager_train', MockIterator,
+        resnet50_test_util.device_and_data_format(), defun=False)
+
+  def benchmark_eager_train_async(self):
+    self._benchmark_eager_train(
+        'eager_train_async',
+        MockIterator,
+        resnet50_test_util.device_and_data_format(),
+        defun=False,
+        execution_mode=context.ASYNC)
+
+  def benchmark_eager_train_with_defun(self):
+    self._benchmark_eager_train(
+        'eager_train_with_defun', MockIterator,
+        resnet50_test_util.device_and_data_format(), defun=True)
+
+  def benchmark_eager_train_datasets(self):
+
+    def make_iterator(tensors):
+      with tf.device('/device:CPU:0'):
+        ds = tf.data.Dataset.from_tensors(tensors).repeat()
+      return iter(ds)
+
+    self._benchmark_eager_train(
+        'eager_train_dataset',
+        make_iterator,
+        resnet50_test_util.device_and_data_format(),
+        defun=False)
+
+  def benchmark_eager_train_datasets_with_defun(self):
+
+    def make_iterator(tensors):
+      with tf.device('/device:CPU:0'):
+        ds = tf.data.Dataset.from_tensors(tensors).repeat()
+      return iter(ds)
+
+    self._benchmark_eager_train(
+        'eager_train_dataset_with_defun', make_iterator,
+        resnet50_test_util.device_and_data_format(), defun=True)
+
+
+if __name__ == '__main__':
+  tf.compat.v1.enable_eager_execution()
+  tf.test.main()

tensorflow/python/eager/benchmarks/resnet50/resnet50_test_util.py

@@ -0,0 +1,55 @@
+# Copyright 2017 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 and benchmarks for the ResNet50 model, executed eagerly."""
+
+import time
+
+import tensorflow as tf
+
+
+def device_and_data_format():
+  if tf.config.list_physical_devices('GPU'):
+    return ('/gpu:0', 'channels_first')
+  return ('/cpu:0', 'channels_last')
+
+
+def random_batch(batch_size, data_format, seed=None):
+  """Create synthetic resnet50 images and labels for testing."""
+  if seed:
+    tf.random.set_seed(seed)
+  shape = (3, 224, 224) if data_format == 'channels_first' else (224, 224, 3)
+  shape = (batch_size,) + shape
+
+  num_classes = 1000
+  images = tf.random.uniform(shape)
+  labels = tf.random.uniform([batch_size],
+                             minval=0,
+                             maxval=num_classes,
+                             dtype=tf.int32)
+  one_hot = tf.one_hot(labels, num_classes)
+
+  return images, one_hot
+
+
+def report(benchmark, label, start, num_iters, device, batch_size, data_format,
+           num_replicas=1):
+  avg_time = (time.time() - start) / num_iters
+  dev = tf.DeviceSpec.from_string(device).device_type.lower()
+  replica_str = '' if num_replicas == 1 else 'replicas_%d_' % num_replicas
+  name = '%s_%s_batch_%d_%s%s' % (label, dev, batch_size,
+                                  replica_str, data_format)
+  extras = {'examples_per_sec': (num_replicas * batch_size) / avg_time}
+  benchmark.report_benchmark(
+      iters=num_iters, wall_time=avg_time, name=name, extras=extras)

tensorflow/python/eager/benchmarks_test_base.py

@@ -0,0 +1,77 @@
+# Copyright 2017 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.
+# ==============================================================================
+r"""Benchmark base to run and report benchmark results."""
+
+import os
+import uuid
+
+from tensorflow.python.eager import test
+from tensorflow.python.platform import flags
+from tensorflow.python.profiler import profiler_v2 as profiler
+
+flags.DEFINE_bool("xprof", False, "Run and report benchmarks with xprof on")
+flags.DEFINE_string("logdir", "/tmp/xprof/", "Directory to store xprof data")
+
+
+class MicroBenchmarksBase(test.Benchmark):
+  """Run and report benchmark results.
+
+  The first run is without any profilng.
+  Second run is with xprof and python trace. Third run is with xprof without
+  python trace. Note: xprof runs are with fewer iterations.
+  """
+
+  def run_with_xprof(self, enable_python_trace, run_benchmark, func,
+                     num_iters_xprof, execution_mode, suid):
+    if enable_python_trace:
+      options = profiler.ProfilerOptions(python_tracer_level=1)
+      logdir = os.path.join(flags.FLAGS.logdir, suid + "_with_python")
+    else:
+      options = profiler.ProfilerOptions(python_tracer_level=0)
+      logdir = os.path.join(flags.FLAGS.logdir, suid)
+    with profiler.Profile(logdir, options):
+      total_time = run_benchmark(func, num_iters_xprof, execution_mode)
+    us_per_example = float("{0:.3f}".format(total_time * 1e6 / num_iters_xprof))
+    return logdir, us_per_example
+
+  def run_report(self, run_benchmark, func, num_iters, execution_mode=None):
+    """Run and report benchmark results."""
+    total_time = run_benchmark(func, num_iters, execution_mode)
+    mean_us = total_time * 1e6 / num_iters
+    extras = {
+        "examples_per_sec": float("{0:.3f}".format(num_iters / total_time)),
+        "us_per_example": float("{0:.3f}".format(total_time * 1e6 / num_iters))
+    }
+
+    if flags.FLAGS.xprof:
+      suid = str(uuid.uuid4())
+      # Re-run with xprof and python trace.
+      num_iters_xprof = min(100, num_iters)
+      xprof_link, us_per_example = self.run_with_xprof(True, run_benchmark,
+                                                       func, num_iters_xprof,
+                                                       execution_mode, suid)
+      extras["xprof link with python trace"] = xprof_link
+      extras["us_per_example with xprof and python"] = us_per_example
+
+      # Re-run with xprof but no python trace.
+      xprof_link, us_per_example = self.run_with_xprof(False, run_benchmark,
+                                                       func, num_iters_xprof,
+                                                       execution_mode, suid)
+      extras["xprof link"] = xprof_link
+      extras["us_per_example with xprof"] = us_per_example
+
+    benchmark_name = self._get_benchmark_name()
+    self.report_benchmark(
+        iters=num_iters, wall_time=mean_us, extras=extras, name=benchmark_name)

tensorflow/python/eager/cancellation.py

@@ -0,0 +1,83 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Cancellation support for eager execution."""
+
+import functools
+import threading
+from typing import Any, Callable, Optional
+
+from tensorflow.python import pywrap_tfe
+
+
+class CancellationManager:
+  """A mechanism for cancelling blocking computation."""
+
+  __slots__ = ["_impl"]
+
+  def __init__(self):
+    self._impl = pywrap_tfe.TFE_NewCancellationManager()
+
+  @property
+  def is_cancelled(self):
+    """Returns `True` if `CancellationManager.start_cancel` has been called."""
+    return pywrap_tfe.TFE_CancellationManagerIsCancelled(self._impl)
+
+  def start_cancel(self):
+    """Cancels blocking operations that have been registered with this object."""
+    pywrap_tfe.TFE_CancellationManagerStartCancel(self._impl)
+
+  def get_cancelable_function(
+      self, concrete_function: Callable[..., Any]
+  ) -> Callable[..., Any]:
+    """Wraps a ConcreteFunction to execute within this cancellation manager's context.
+
+    Args:
+      concrete_function: The ConcreteFunction to wrap.
+
+    Returns:
+      A callable that executes the concrete_function under this cancellation
+      context.
+    """
+
+    @functools.wraps(concrete_function)
+    def cancellable(*args, **kwargs):
+      with CancellationManagerContext(self):
+        return concrete_function(*args, **kwargs)
+
+    return cancellable
+
+
+_active_context = threading.local()
+
+
+def context() -> Optional[CancellationManager]:
+  """Returns the CancellationManager active in the current thread, or None."""
+  stack = getattr(_active_context, "manager_stack", None)
+  return stack[-1] if stack else None
+
+
+class CancellationManagerContext:
+  """A Python context for wrapping a cancellable ConcreteFunction."""
+
+  def __init__(self, cancellation_manager):
+    self._cancellation_manager = cancellation_manager
+
+  def __enter__(self):
+    if not hasattr(_active_context, "manager_stack"):
+      _active_context.manager_stack = []
+    _active_context.manager_stack.append(self._cancellation_manager)
+
+  def __exit__(self, exc_type, exc_value, exc_tb):
+    _active_context.manager_stack.pop()

tensorflow/python/eager/cancellation_test.py

@@ -0,0 +1,77 @@
+# Copyright 2019 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.
+# ==============================================================================
+
+import threading
+from tensorflow.python.eager import cancellation
+from tensorflow.python.platform import test
+
+
+class CancellationTest(test.TestCase):
+
+  def test_start_cancel(self):
+    manager = cancellation.CancellationManager()
+    self.assertFalse(manager.is_cancelled)
+    manager.start_cancel()
+    self.assertTrue(manager.is_cancelled)
+
+  def test_threading(self):
+    manager1 = cancellation.CancellationManager()
+    manager2 = cancellation.CancellationManager()
+    thread_exceptions = []
+
+    def thread_fn():
+      try:
+        with cancellation.CancellationManagerContext(manager2):
+          self.assertEqual(cancellation.context(), manager2)
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        thread_exceptions.append(e)
+
+    with cancellation.CancellationManagerContext(manager1):
+      self.assertEqual(cancellation.context(), manager1)
+      t = threading.Thread(target=thread_fn)
+      t.start()
+      t.join()
+      for exc in thread_exceptions:
+        raise exc
+      self.assertEqual(cancellation.context(), manager1)
+
+  def test_nested_context(self):
+    manager1 = cancellation.CancellationManager()
+    manager2 = cancellation.CancellationManager()
+
+    self.assertIsNone(cancellation.context())
+    with cancellation.CancellationManagerContext(manager1):
+      self.assertEqual(cancellation.context(), manager1)
+      with cancellation.CancellationManagerContext(manager2):
+        self.assertEqual(cancellation.context(), manager2)
+      self.assertEqual(cancellation.context(), manager1)
+    self.assertIsNone(cancellation.context())
+
+  def test_get_cancelable_function(self):
+    manager = cancellation.CancellationManager()
+
+    def my_fn(x):
+      """My docstring."""
+      self.assertEqual(cancellation.context(), manager)
+      return x + 1
+
+    cancelable_fn = manager.get_cancelable_function(my_fn)
+    self.assertEqual(cancelable_fn.__name__, "my_fn")
+    self.assertEqual(cancelable_fn.__doc__, "My docstring.")
+    self.assertEqual(cancelable_fn(5), 6)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/context_cross_platform_gpu_test.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ContextCrossPlatformGpuTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      [(f'_{stage}', stage) for stage in ['hlo', 'hlo_serialized']]
+  )
+  def testGetCompilerIrOnTpuPlatform(self, stage):
+    @def_function.function(jit_compile=True)
+    def test_func(x):
+      return 2 * x
+
+    a = array_ops.ones((1000, 1000))  # 4 * 1000 * 1000 in bytes
+    result = test_func.experimental_get_compiler_ir(a)(
+        stage=stage, platform_name='TPU'
+    )
+    self.assertNotEmpty(result)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/context_cross_platform_tpu_test.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+class ContextCrossPlatformTpuTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      [(f'_{stage}', stage) for stage in ['hlo', 'hlo_serialized']]
+  )
+  def testGetCompilerIrOnGpuPlatform(self, stage):
+    @def_function.function(jit_compile=True)
+    def test_func(x):
+      return 2.0 * x
+
+    a = constant_op.constant(1.0)
+    result = test_func.experimental_get_compiler_ir(a)(
+        stage=stage, platform_name='GPU'
+    )
+    self.assertNotEmpty(result)
+
+  @parameterized.named_parameters([
+      (f'_{stage}', stage)
+      for stage in [
+          'optimized_hlo',
+          'optimized_hlo_serialized',
+          'optimized_hlo_proto_serialized',
+          'optimized_hlo_dot',
+      ]
+  ])
+  def testGetCompilerIrOnGpuPlatformOptimizedHlo(self, stage):
+    @def_function.function(jit_compile=True)
+    def test_func(x):
+      return 2.0 * x
+
+    a = constant_op.constant(1.0)
+    with self.assertRaisesRegex(
+        ValueError,
+        'GetCompilerIr with requested stage is not supported on this device',
+    ):
+      _ = test_func.experimental_get_compiler_ir(a)(
+          stage=stage, platform_name='GPU'
+      )
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/context_test.py

@@ -0,0 +1,217 @@
+# Copyright 2019 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.
+# ==============================================================================
+
+import weakref
+
+from absl.testing import parameterized
+import numpy as np
+
+from xla.service import hlo_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ContextTest(test.TestCase, parameterized.TestCase):
+
+  def testSetGlobalSeed(self):
+    c = context.Context()
+    c._set_global_seed(123)
+    for t in [np.int32, np.int64, np.uint32, np.uint64]:
+      c._set_global_seed(t(123))
+      c._set_global_seed(np.array(123, dtype=t))
+      c._set_global_seed(ops.convert_to_tensor(123, dtype=t))
+
+  def testContextIsDestroyedAfterTensors(self):
+    # Create a new context
+    new_context = context.Context()
+    weak_c = weakref.ref(new_context)
+    new_context.ensure_initialized()
+
+    # Create a tensor with the new context as default.
+    # Make sure to restore the original context.
+    original_context = context.context()
+    try:
+      context._set_context(new_context)
+      # Use a 2D tensor so that it is not cached.
+      tensor1 = constant_op.constant([[3.]])
+      # Produce a tensor as an operation output. This uses a different code path
+      # from tensors created from Python.
+      tensor2 = tensor1 * tensor1
+      context._set_context(original_context)
+    except:
+      context._set_context(original_context)
+      raise
+
+    # Deleting our context reference should not delete the underlying object.
+    del new_context
+    self.assertIsNot(weak_c(), None)
+
+    # Deleting the first tensor should not delete the context since there is
+    # another tensor.
+    del tensor1
+    self.assertIsNot(weak_c(), None)
+
+    # Deleting the last tensor should result in deleting its context.
+    del tensor2
+    self.assertIs(weak_c(), None)
+
+  def testSimpleGraphCollection(self):
+
+    @def_function.function
+    def f(x):
+      with ops.device('CPU:0'):
+        return x + constant_op.constant(1.)
+
+    with context.collect_graphs() as graphs:
+      with ops.device('CPU:0'):
+        x = constant_op.constant(1.)
+      f(x)
+
+    self.assertLen(graphs, 1)
+    graph, = graphs
+    self.assertIn('CPU:0', graph.node[1].device)
+
+  @test_util.disable_tfrt(
+      'b/171600738: tfrt does not support exporting post-optimization graph')
+  def testGraphCollectionAfterDevicePlacement(self):
+
+    @def_function.function
+    def f(x):
+      return x + constant_op.constant(1.)
+
+    with context.collect_graphs() as graphs:
+      with ops.device('CPU:0'):
+        f(constant_op.constant(1.))
+
+    self.assertLen(graphs, 1)
+    graph, = graphs
+    self.assertIn('CPU:0', graph.node[0].device)
+
+  def testGetFunctionDef(self):
+
+    @def_function.function
+    def f():
+      return constant_op.constant(1.)
+
+    concrete = f.get_concrete_function()
+    function_def = context.get_function_def(concrete.name)
+
+    self.assertIsNot(function_def, None)
+
+    found_const_node = False
+    for node_def in function_def.node_def:
+      if node_def.op == 'Const':
+        found_const_node = True
+        break
+    self.assertTrue(found_const_node)
+
+    with self.assertRaises(errors.NotFoundError):
+      _ = context.get_function_def('this_should_not_be_found')
+
+  @test_util.run_gpu_only
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfo(self):
+    array_ops.zeros([10])  # Allocate some memory on the GPU.
+    self.assertGreater(context.context().get_memory_info('GPU:0')['current'], 0)
+
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfoCPU(self):
+    if test_util.IsMklEnabled():
+      # TODO(gzmkl) work with Google team to address design issue in allocator.h
+      self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
+
+    with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
+      context.context().get_memory_info('CPU:0')
+
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfoUnknownDevice(self):
+    with self.assertRaisesRegex(ValueError, 'No matching devices found'):
+      context.context().get_memory_info('unknown_device:0')
+
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfoUnparsableDevice(self):
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU:')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU:CPU')
+
+  def testListFunctionNames(self):
+
+    @def_function.function
+    def f():
+      return constant_op.constant(1.)
+
+    concrete = f.get_concrete_function()
+    self.assertIn(concrete.name.decode(),
+                  context.context().list_function_names())
+
+  def testSetLogicalDeviceAfterContextInitialization(self):
+    ctx = context.Context()
+    ctx.set_logical_cpu_devices(4)
+    self.assertIs(len(ctx.list_logical_devices('CPU')), 4)
+
+    # Cannot set logical device twice.
+    with self.assertRaisesRegex(RuntimeError, 'Virtual CPUs already set'):
+      ctx.set_logical_cpu_devices(8)
+
+  def testSetLogicalLocalCpuDevice(self):
+    ctx = context.Context()
+
+    # Manually add a remote CPU device into logical device list.
+    ctx._logical_devices = []  # pylint: disable=protected-access
+    dev = context.LogicalDevice(name='/job:worker/replica:0/task:1',
+                                device_type='CPU')
+    ctx._logical_devices.append(dev)  # pylint: disable=protected-access
+    self.assertIs(len(ctx.list_logical_devices('CPU')), 1)
+
+    # This would pass the check since the previously added device is not local.
+    ctx.set_logical_cpu_devices(4)
+    # Logical device list would be overwritten after initialization.
+    self.assertIs(len(ctx.list_logical_devices('CPU')), 4)
+
+  @parameterized.named_parameters([(f'_{stage}', stage) for stage in [
+      'hlo', 'hlo_serialized', 'optimized_hlo', 'optimized_hlo_serialized',
+      'optimized_hlo_proto_serialized', 'optimized_hlo_dot'
+  ]])
+  def testGetCompilerIr(self, stage):
+
+    @def_function.function(jit_compile=True)
+    def test_func(x):
+      return 2 * x
+
+    a = array_ops.ones((1000, 1000))  # 4 * 1000 * 1000 in bytes
+    result = test_func.experimental_get_compiler_ir(a)(stage=stage)
+    self.assertNotEmpty(result)
+    if stage == 'optimized_hlo_proto_serialized':
+      hlo_proto = hlo_pb2.HloProto.FromString(result)
+      allocations = hlo_proto.buffer_assignment.buffer_allocations
+      buffer_size = sum(
+          getattr(allocation, 'size') for allocation in allocations)
+      # The sizes of input and output are both 4 * 1000 * 1000 in bytes.
+      self.assertGreaterEqual(buffer_size, 2 * 4 * 1000 * 1000)
+      self.assertLess(buffer_size, 4 * 4 * 1000 * 1000)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/context_tpu_platform_test.py

@@ -0,0 +1,47 @@
+# 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class ContextTpuPlatformTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      [(f'_{stage}', stage) for stage in ['hlo', 'hlo_serialized']]
+  )
+  def testGetCompilerIrWithVariable(self, stage):
+    with ops.device('TPU:0'):
+      v = variables.Variable(1.0)
+
+    @def_function.function(jit_compile=True)
+    def test_func(x):
+      return x * v
+
+    a = constant_op.constant(1.0)
+    result = test_func.experimental_get_compiler_ir(a)(
+        stage=stage, platform_name='TPU'
+    )
+    self.assertNotEmpty(result)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/core.py

@@ -0,0 +1,78 @@
+# Copyright 2017 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 API for TensorFlow's "Eager" mode of execution."""
+
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import tf_logging as logging
+
+# Trace of execution and memory usage.
+_active_trace = None
+
+
+def _status_to_exception(status):
+  try:
+    error_class = errors.exception_type_from_error_code(status.code)
+    e = error_class(None, None, status.message, status.payloads)
+    logging.error_log("%s: %s" % (e.__class__.__name__, e))
+    return e
+  except KeyError:
+    e = errors.UnknownError(
+        None, None, status.message, status.code, status.payloads
+    )
+    logging.error_log("%s: %s" % (e.__class__.__name__, e))
+    return e
+
+
+class _NotOkStatusException(Exception):
+  """Exception class to handle not ok Status."""
+
+  def __init__(self, message, code, payloads):
+    super(_NotOkStatusException, self).__init__()
+    self.message = message
+    self.code = code
+    self.payloads = payloads
+
+  def __str__(self):
+    e = _status_to_exception(self)
+    return "%s: %s" % (e.__class__.__name__, e)
+
+
+pywrap_tfe.TFE_Py_RegisterExceptionClass(_NotOkStatusException)
+
+
+class _FallbackException(Exception):
+  """Exception class to handle fallback from the fastpath.
+
+  The fastpath that we refer to here is the one implemented to reduce per-op
+  overheads (TFE_Py_FastPathExecute_C). If the conditions for executing the op
+  on the fastpath are not met, we fallback to a safer (and more complete)
+  slowpath, and this Exception is raised to signal that transition.
+  """
+  pass
+
+
+class _SymbolicException(Exception):
+  """Exception class to handle use of symbolic tensors when executing eagerly.
+
+  `keras.Input()` creates symbolic tensors (in a FuncGraph managed by the
+  Keras backend) while in eager execution. This exception is used to
+  identify this case (raised in `convert_to_tensor` cause generated functions
+  for ops to construct graphs instead of executing the kernel).
+  """
+  pass
+
+
+pywrap_tfe.TFE_Py_RegisterFallbackExceptionClass(_FallbackException)

tensorflow/python/eager/custom_device_test.py

@@ -0,0 +1,59 @@
+# 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.
+# ==============================================================================
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import custom_device_testutil
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+class CustomDeviceTest(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    context._reset_context()
+
+  def testRegisterCustomDevice(self):
+    device_name = '/job:localhost/replica:0/task:0/device:CUSTOM:0'
+    device, device_info, arrived_flag, executed_flag = (
+        custom_device_testutil.GetLoggingDeviceCapsules(device_name))
+    context.register_custom_device(device, device_name, device_info)
+    self.assertFalse(custom_device_testutil.FlagValue(arrived_flag))
+    self.assertFalse(custom_device_testutil.FlagValue(executed_flag))
+    with ops.device(device_name):
+      x = constant_op.constant(1.)
+      y = x * constant_op.constant(2.)
+    self.assertTrue(custom_device_testutil.FlagValue(executed_flag))
+    # There was no copy onto the device. Actually I'm not sure how to trigger
+    # that from Python.
+    self.assertFalse(custom_device_testutil.FlagValue(arrived_flag))
+    with self.assertRaisesRegex(errors.InternalError, 'Trying to copy'):
+      y.numpy()
+
+  def testIsCustomDevice(self):
+    device_name = '/job:localhost/replica:0/task:0/device:CUSTOM:0'
+    device, device_info, _, _ = (
+        custom_device_testutil.GetLoggingDeviceCapsules(device_name))
+    context.register_custom_device(device, device_name, device_info)
+
+    self.assertTrue(context.is_custom_device(device_name))
+    self.assertFalse(context.is_custom_device('cpu:0'))
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/custom_device_testutil.cc

@@ -0,0 +1,69 @@
+/* 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/eager/custom_device_testutil.h"
+
+#include "Python.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/c/c_api.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/c_api_experimental.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;
+
+void CallDelete_Device(PyObject* capsule) {
+  delete reinterpret_cast<TFE_CustomDevice*>(
+      PyCapsule_GetPointer(capsule, "TFE_CustomDevice"));
+}
+
+void CallDelete_DeviceInfo(PyObject* capsule) {
+  PyErr_SetString(PyExc_AssertionError,
+                  "Capsule should be consumed by TFE_Py_RegisterCustomDevice");
+  DeleteLoggingDevice(
+      PyCapsule_GetPointer(capsule, "TFE_CustomDevice_DeviceInfo"));
+}
+
+PYBIND11_MODULE(custom_device_testutil, m) {
+  m.def("GetLoggingDeviceCapsules", [](const char* name) {
+    bool* arrived_flag = new bool;
+    bool* executed_flag = new bool;
+    *arrived_flag = false;
+    *executed_flag = false;
+    tensorflow::Safe_PyObjectPtr arrived_capsule(
+        PyCapsule_New(arrived_flag, "flag", nullptr));
+    tensorflow::Safe_PyObjectPtr executed_capsule(
+        PyCapsule_New(executed_flag, "flag", nullptr));
+    TFE_CustomDevice* device;
+    void* device_info;
+    AllocateLoggingDevice(name, arrived_flag, executed_flag, &device,
+                          &device_info);
+    tensorflow::Safe_PyObjectPtr device_capsule(
+        PyCapsule_New(device, "TFE_CustomDevice", &CallDelete_Device));
+    tensorflow::Safe_PyObjectPtr device_info_capsule(PyCapsule_New(
+        device_info, "TFE_CustomDevice_DeviceInfo", &CallDelete_DeviceInfo));
+    return tensorflow::PyoOrThrow(
+        PyTuple_Pack(4, device_capsule.get(), device_info_capsule.get(),
+                     arrived_capsule.get(), executed_capsule.get()));
+  });
+  m.def("FlagValue", [](py::capsule flag_capsule) {
+    bool* flag = reinterpret_cast<bool*>(
+        PyCapsule_GetPointer(flag_capsule.ptr(), "flag"));
+    if (PyErr_Occurred()) throw py::error_already_set();
+    return *flag;
+  });
+}

tensorflow/python/eager/custom_device_testutil.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 FlagValue(arg0) -> bool: ...
+def GetLoggingDeviceCapsules(arg0: str) -> object: ...

tensorflow/python/eager/def_function.py

@@ -0,0 +1,28 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Supports old symbols supplied by this file while the code is refactored."""
+
+# pylint:disable=unused-import,g-bad-import-order
+
+# Config Options
+from tensorflow.python.eager.polymorphic_function.eager_function_run import run_functions_eagerly
+from tensorflow.python.eager.polymorphic_function.eager_function_run import functions_run_eagerly
+
+# tf.function Classes
+from tensorflow.python.eager.polymorphic_function.polymorphic_function import Function
+from tensorflow.python.eager.polymorphic_function.polymorphic_function import function
+
+# Private attributes
+from tensorflow.python.eager.polymorphic_function.polymorphic_function import _tf_function_counter

tensorflow/python/eager/device_placement_test.py

@@ -0,0 +1,298 @@
+# Copyright 2019 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 device placement."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+
+
+class SoftDevicePlacementTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(SoftDevicePlacementTest, self).setUp()
+    context._reset_context()
+    context.ensure_initialized()
+    config.set_soft_device_placement(enabled=True)
+    context.context().log_device_placement = True
+
+  @test_util.run_gpu_only
+  def testDefaultPlacement(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    c = a + b
+    with ops.device('CPU'):
+      d = a + b
+    self.assertIn('GPU', c.device)
+    self.assertIn('CPU', d.device)
+
+  @test_util.run_gpu_only
+  def testUnsupportedDevice(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    s = constant_op.constant(list('hello world'))
+    with ops.device('GPU:0'):
+      c = a + b
+      t = s[a]
+    self.assertIn('GPU:0', c.device)
+    self.assertIn('CPU', t.device)
+
+  @test_util.run_gpu_only
+  def testUnknownDevice(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    with ops.device('GPU:42'):
+      c = a + b
+    self.assertIn('GPU:0', c.device)
+
+  def testNoGpu(self):
+    if test_util.is_gpu_available():
+      # CPU only test.
+      return
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    c = a + b
+    with ops.device('GPU'):
+      d = a + b
+    self.assertIn('CPU', c.device)
+    self.assertIn('CPU', d.device)
+
+  @test_util.run_gpu_only
+  def testSoftPlacedGPU(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    with ops.device('GPU:110'):
+      c = a + b
+    self.assertIn('GPU:0', c.device)
+
+  @test_util.run_gpu_only
+  def testNestedDeviceScope(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    with ops.device('CPU:0'):
+      with ops.device('GPU:42'):
+        c = a + b
+    # We don't support nested device placement right now.
+    self.assertIn('GPU:0', c.device)
+
+  @parameterized.named_parameters(('float', 1.0, None),
+                                  ('int32', [1], dtypes.int32),
+                                  ('string', ['a'], None))
+  def testSoftPlacedCPUConstant(self, value, dtype):
+    if test_util.is_gpu_available():
+      self.skipTest('CPU only test')
+    with ops.device('GPU:0'):
+      a = constant_op.constant(value, dtype=dtype)
+    self.assertIn('CPU:0', a.device)
+    self.assertIn('CPU:0', a.backing_device)
+
+  @parameterized.named_parameters(
+      ('float', [1.0, 2.0], None, 'GPU:0'),
+      ## TODO(b/179035075) enable when int32 numeric tensors are not confused
+      ## with shape tensors
+      # ('int32', [1, 2], dtypes.int32, 'GPU:0'),
+      ('int64', [1, 2], dtypes.int64, 'GPU:0'))
+  @test_util.run_gpu_only
+  def testSoftPlacedNumericTensors(self, value, dtype, expect):
+    with ops.device('GPU:0'):
+      a = math_ops.add(constant_op.constant(value, dtype=dtype),
+                       constant_op.constant(value, dtype=dtype))
+    self.assertIn(expect, a.backing_device)
+
+  @test_util.run_gpu_only
+  def testSoftPlacedShapeTensor(self):
+    with ops.device('GPU:0'):
+      t = constant_op.constant([[1, 2], [3, 4]])
+      a = math_ops.add(array_ops.shape(t),
+                       constant_op.constant([10, 20], dtype=dtypes.int32))
+    # Shape computations remain on CPU
+    self.assertIn('CPU:0', a.backing_device)
+
+  def testPlacedToDeviceInFunction(self):
+
+    @def_function.function
+    def f():
+      a = random_ops.random_uniform([32, 32])
+      return math_ops.matmul(a, a)
+
+    gpus = config.list_physical_devices('GPU')
+    if not gpus:
+      self.assertIn('CPU:0', f().device)
+    else:
+      self.assertIn('GPU:0', f().device)
+
+  @test_util.disable_tfrt('b/173726713: Support properly inserting device at '
+                          'tf_to_corert lowering.')
+  def testUnknownDeviceInFunction(self):
+
+    @def_function.function
+    def f():
+      with ops.device('GPU:42'):
+        # With placer, the unknown GPU:42 will be replaced with GPU:0.
+        a = constant_op.constant(1) + constant_op.constant(2)
+      return a + constant_op.constant(2)
+
+    gpus = config.list_physical_devices('GPU')
+    if not gpus:
+      self.assertIn('CPU:0', f().device)
+    else:
+      self.assertIn('GPU:0', f().device)
+
+
+class HardDevicePlacementTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(HardDevicePlacementTest, self).setUp()
+    context._reset_context()
+    config.set_soft_device_placement(enabled=False)
+    context.context().log_device_placement = True
+    cpus = context.context().list_physical_devices('CPU')
+    # Set 2 virtual CPUs
+    context.context().set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+    self.assertEqual(config.get_soft_device_placement(), False)
+    self.assertEqual(context.context().soft_device_placement, False)
+
+  @test_util.run_gpu_only
+  def testIdentityCanCopy(self):
+    config.set_device_policy('explicit')
+    with ops.device('CPU:0'):
+      x = constant_op.constant(1.0)
+      self.assertIn('CPU:0', x.device)
+      self.assertIn('CPU:0', x.backing_device)
+    with ops.device('GPU:0'):
+      y = array_ops.identity(x)
+      self.assertIn('GPU:0', y.device)
+      self.assertIn('GPU:0', y.backing_device)
+
+  @test_util.run_gpu_only
+  def testSimpleConstantsExplicitGPU(self):
+    config.set_device_policy('explicit')
+    with ops.device('GPU:0'):
+      self.assertAllClose(1., array_ops.ones([]))
+      self.assertAllClose(0., array_ops.zeros([]))
+      self.assertAllClose([1.], array_ops.fill([1], 1.))
+
+  def testSimpleConstantsExplicitCPU(self):
+    config.set_device_policy('explicit')
+    with ops.device('CPU:1'):
+      self.assertAllClose(1., array_ops.ones([]))
+      self.assertAllClose(0., array_ops.zeros([]))
+      self.assertAllClose([1.], array_ops.fill([1], 1.))
+      self.assertAllClose(2., constant_op.constant(1.) * 2.)
+
+  @parameterized.named_parameters(
+      ('float_cpu0', 'CPU:0', 1.0, None),
+      ('int32_cpu0', 'CPU:0', [1], dtypes.int32),
+      ('string_cpu0', 'CPU:0', ['a'], None),
+      ('float_cpu1', 'CPU:1', 1.0, None),
+      ('int32_cpu1', 'CPU:1', [1], dtypes.int32),
+      ('string_cpu1', 'CPU:1', ['a'], None),
+  )
+  def testHardPlacedCPUConstant(self, device, value, dtype):
+    with ops.device(device):
+      a = constant_op.constant(value, dtype=dtype)
+      self.assertIn(device, a.device)
+
+  @parameterized.named_parameters(
+      ('float', 'GPU:0', 1.0, None),
+      ('int32', 'GPU:0', [1], dtypes.int32),
+      ('string', 'GPU:0', ['a'], None),
+  )
+  def testHardPlacedGPUConstant(self, device, value, dtype):
+    if not test_util.is_gpu_available():
+      self.skipTest('Test requires a GPU')
+    with ops.device(device):
+      a = constant_op.constant(value, dtype=dtype)
+      self.assertIn(device, a.device)
+      if a.dtype == dtypes.float32:
+        self.assertIn(device, a.backing_device)
+
+
+class ClusterPlacementTest(test.TestCase):
+
+  def setUp(self):
+    super(ClusterPlacementTest, self).setUp()
+    context._reset_context()
+    config.set_soft_device_placement(enabled=True)
+    context.context().log_device_placement = True
+    workers, _ = test_util.create_local_cluster(2, 0)
+    remote.connect_to_remote_host([workers[0].target, workers[1].target])
+
+  @test_util.disable_tfrt('remote host not supported yet.')
+  def testNotFullySpecifiedTask(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    with ops.device('/job:worker'):
+      c = a + b
+    self.assertIn('/job:worker/replica:0/task:0', c.device)
+
+  @test_util.disable_tfrt('remote host not supported yet.')
+  def testRemoteUnknownDevice(self):
+    a = constant_op.constant(1)
+    b = constant_op.constant(2)
+    # Right now we don't support soft device place on remote worker.
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      with ops.device('/job:worker/replica:0/task:0/device:GPU:42'):
+        c = a + b
+        del c
+      self.assertIn('unknown device', cm.exception.message)
+
+  @test_util.disable_tfrt('remote host not supported yet.')
+  def testUnknownDeviceInFunctionReturnUnknowDevice(self):
+
+    @def_function.function
+    def f():
+      with ops.device('GPU:42'):
+        return constant_op.constant(1) + constant_op.constant(2)
+
+    gpus = config.list_physical_devices('GPU')
+    if not gpus:
+      self.assertIn('CPU:0', f().device)
+    else:
+      self.assertIn('GPU:0', f().device)
+
+  @test_util.disable_tfrt('remote host not supported yet.')
+  def testUnknownDeviceInFunction(self):
+
+    @def_function.function
+    def f():
+      with ops.device('GPU:42'):
+        a = constant_op.constant(1) + constant_op.constant(2)
+      return a + constant_op.constant(2)
+
+    gpus = config.list_physical_devices('GPU')
+    if not gpus:
+      self.assertIn('CPU:0', f().device)
+    else:
+      self.assertIn('GPU:0', f().device)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/execute.py

@@ -0,0 +1,329 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Functions called by the generated code to execute an eager-mode op."""
+
+from google.protobuf import text_format
+from tensorflow.core.framework import tensor_pb2
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.eager import core
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_conversion_registry
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.types import core as core_types
+from tensorflow.python.util import compat
+
+
+def quick_execute(op_name, num_outputs, inputs, attrs, ctx, name=None):
+  """Execute a TensorFlow operation.
+
+  Args:
+    op_name: Name of the TensorFlow operation (see REGISTER_OP in C++ code) to
+      execute.
+    num_outputs: The number of outputs of the operation to fetch. (Explicitly
+      provided instead of being inferred for performance reasons).
+    inputs: A list of inputs to the operation. Each entry should be a Tensor, or
+      a value which can be passed to the Tensor constructor to create one.
+    attrs: A tuple with alternating string attr names and attr values for this
+      operation.
+    ctx: The value of context.context().
+    name: Customized name for the operation.
+
+  Returns:
+    List of output Tensor objects. The list is empty if there are no outputs
+
+  Raises:
+    An exception on error.
+  """
+  device_name = ctx.device_name
+  # pylint: disable=protected-access
+  try:
+    ctx.ensure_initialized()
+    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
+                                        inputs, attrs, num_outputs)
+  except core._NotOkStatusException as e:
+    if name is not None:
+      e.message += " name: " + name
+    raise core._status_to_exception(e) from None
+  except TypeError as e:
+    keras_symbolic_tensors = [x for x in inputs if _is_keras_symbolic_tensor(x)]
+    if keras_symbolic_tensors:
+      raise core._SymbolicException(
+          "Inputs to eager execution function cannot be Keras symbolic "
+          "tensors, but found {}".format(keras_symbolic_tensors))
+    raise e
+  # pylint: enable=protected-access
+  return tensors
+
+
+def execute_with_cancellation(op_name,
+                              num_outputs,
+                              inputs,
+                              attrs,
+                              ctx,
+                              cancellation_manager,
+                              name=None):
+  """Execute a TensorFlow operation.
+
+  Args:
+    op_name: Name of the TensorFlow operation (see REGISTER_OP in C++ code) to
+      execute.
+    num_outputs: The number of outputs of the operation to fetch. (Explicitly
+      provided instead of being inferred for performance reasons).
+    inputs: A list of inputs to the operation. Each entry should be a Tensor, or
+      a value which can be passed to the Tensor constructor to create one.
+    attrs: A tuple with alternating string attr names and attr values for this
+      operation.
+    ctx: The value of context.context().
+    cancellation_manager: a `CancellationManager` object that can be used to
+      cancel the operation.
+    name: Customized name for the operation.
+
+  Returns:
+    List of output Tensor objects. The list is empty if there are no outputs
+
+  Raises:
+    An exception on error.
+  """
+  device_name = ctx.device_name
+  # pylint: disable=protected-access
+  try:
+    ctx.ensure_initialized()
+    tensors = pywrap_tfe.TFE_Py_ExecuteCancelable(ctx._handle, device_name,
+                                                  op_name, inputs, attrs,
+                                                  cancellation_manager._impl,
+                                                  num_outputs)
+  except core._NotOkStatusException as e:
+    if name is not None:
+      e.message += " name: " + name
+    raise core._status_to_exception(e) from None
+  except TypeError as e:
+    keras_symbolic_tensors = [x for x in inputs if _is_keras_symbolic_tensor(x)]
+    if keras_symbolic_tensors:
+      raise core._SymbolicException(
+          "Inputs to eager execution function cannot be Keras symbolic "
+          "tensors, but found {}".format(keras_symbolic_tensors))
+    raise e
+  # pylint: enable=protected-access
+  return tensors
+
+
+def execute_with_callbacks(op_name, num_outputs, inputs, attrs, ctx, name=None):
+  """Monkey-patch to execute to enable execution callbacks."""
+  tensors = quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
+  for callback in ctx.op_callbacks:
+    callback(op_name, tuple(inputs), attrs, tensors, name)
+
+  return tensors
+
+
+execute = quick_execute
+
+
+def must_record_gradient():
+  """Import backprop if you want gradients recorded."""
+  return False
+
+
+def record_gradient(unused_op_name, unused_inputs, unused_attrs,
+                    unused_outputs):
+  """Import backprop if you want gradients recorded."""
+  pass
+
+
+def make_float(v, arg_name):
+  if not isinstance(v, compat.real_types):
+    raise TypeError("Expected float for argument '%s' not %s." %
+                    (arg_name, repr(v)))
+  return float(v)
+
+
+def make_int(v, arg_name):
+  if isinstance(v, str):
+    raise TypeError("Expected int for argument '%s' not %s." %
+                    (arg_name, repr(v)))
+  try:
+    return int(v)
+  except (ValueError, TypeError):
+    raise TypeError("Expected int for argument '%s' not %s." %
+                    (arg_name, repr(v)))
+
+
+def make_str(v, arg_name):
+  if not isinstance(v, compat.bytes_or_text_types):
+    raise TypeError("Expected string for argument '%s' not %s." %
+                    (arg_name, repr(v)))
+  return compat.as_bytes(v)  # Convert unicode strings to bytes.
+
+
+def make_bool(v, arg_name):
+  if not isinstance(v, bool):
+    raise TypeError("Expected bool for argument '%s' not %s." %
+                    (arg_name, repr(v)))
+  return v
+
+
+def make_type(v, arg_name):
+  try:
+    v = dtypes.as_dtype(v).base_dtype
+  except TypeError:
+    raise TypeError("Expected DataType for argument '%s' not %s." %
+                    (arg_name, repr(v)))
+  i = v.as_datatype_enum
+  return i
+
+
+def make_shape(v, arg_name):
+  """Convert v into a list."""
+  # Args:
+  #   v: A TensorShapeProto, a list of ints, or a tensor_shape.TensorShape.
+  #   arg_name: String, for error messages.
+
+  # Returns:
+  #   None if the rank is unknown, otherwise a list of ints (or Nones in the
+  #   position where the dimension is unknown).
+  try:
+    shape = tensor_shape.as_shape(v)
+  except TypeError as e:
+    raise TypeError("Error converting %s to a TensorShape: %s." % (arg_name, e))
+  except ValueError as e:
+    raise ValueError("Error converting %s to a TensorShape: %s." %
+                     (arg_name, e))
+  if shape.ndims is None:
+    return None
+  else:
+    return shape.as_list()
+
+
+def make_tensor(v, arg_name):
+  """Ensure v is a TensorProto."""
+  if isinstance(v, tensor_pb2.TensorProto):
+    return v
+  elif isinstance(v, str):
+    pb = tensor_pb2.TensorProto()
+    text_format.Merge(v, pb)
+    return pb
+  raise TypeError(
+      "Don't know how to convert %s to a TensorProto for argument '%s'." %
+      (repr(v), arg_name))
+
+
+def args_to_matching_eager(l, ctx, allowed_dtypes, default_dtype=None):
+  """Convert sequence `l` to eager same-type Tensors."""
+  del ctx  # Unused
+  if (not l) and (default_dtype is not None):
+    return default_dtype, []  # List is empty; assume default dtype.
+  for x in l:
+    if not isinstance(x, core_types.Value):
+      break
+  else:  # note: intentional for-else
+    return l[0]._datatype_enum(), l  # pylint: disable=protected-access
+
+  # Is some input already a Tensor with a dtype?
+  dtype = None
+  for t in l:
+    if isinstance(t, core_types.Value):
+      dtype = t.dtype
+      break
+
+  if dtype is None:
+    # Infer a dtype based on the first value, and use that dtype for the
+    # remaining values.
+
+    ret = []
+    for t in l:
+      tensor = None
+      # First see if we can get a valid dtype with the default conversion
+      # and see if it matches an allowed dtypes. Some ops like ConcatV2 may
+      # not list allowed dtypes, in which case we should skip this.
+      if dtype is None and allowed_dtypes:
+        tensor = tensor_conversion_registry.convert(t)
+        # If we did not match an allowed dtype, try again with the default
+        # dtype. This could be because we have an empty tensor and thus we
+        # picked the wrong type.
+        if tensor.dtype not in allowed_dtypes:
+          tensor = None
+
+      if tensor is None:
+        tensor = tensor_conversion_registry.convert(
+            t, dtype, preferred_dtype=default_dtype
+        )
+
+      ret.append(tensor)
+      if dtype is None:
+        dtype = tensor.dtype
+  else:
+    ret = [tensor_conversion_registry.convert(t, dtype) for t in l]
+
+  # TODO(slebedev): consider removing this as it leaks a Keras concept.
+  # pylint: disable=protected-access
+  keras_symbolic_tensors = [x for x in ret if _is_keras_symbolic_tensor(x)]
+  if keras_symbolic_tensors:
+    raise core._SymbolicException(
+        "Using symbolic output of a Keras layer during eager execution "
+        "{}".format(keras_symbolic_tensors))
+  # pylint: enable=protected-access
+  return dtype.as_datatype_enum, ret
+
+
+def convert_to_mixed_eager_tensors(values, ctx):
+  del ctx  # Unused
+  v = [tensor_conversion_registry.convert(t) for t in values]
+  types = [t._datatype_enum() for t in v]  # pylint: disable=protected-access
+  return types, v
+
+
+def args_to_mixed_eager_tensors(lists, ctx):
+  """Converts a list of same-length lists of values to eager tensors."""
+  del ctx  # Unused
+  assert len(lists) > 1
+
+  # Generate an error if len(lists[i]) is not the same for all i.
+  lists_ret = [[]]
+  for l in lists[1:]:
+    if len(l) != len(lists[0]):
+      raise ValueError(
+          "Expected list arguments to be the same length: %d != %d (%r vs. %r)."
+          % (len(lists[0]), len(l), lists[0], l))
+    lists_ret.append([])
+
+  # Convert the first element of each list first, then the second element, etc.
+  types = []
+  for i in range(len(lists[0])):
+    dtype = None
+    # If any list has a Tensor, use that dtype
+    for l in lists:
+      if isinstance(l[i], core_types.Value):
+        dtype = l[i].dtype
+        break
+    if dtype is None:
+      # Convert the first one and use its dtype.
+      lists_ret[0].append(tensor_conversion_registry.convert(lists[0][i]))
+      dtype = lists_ret[0][i].dtype
+      for j in range(1, len(lists)):
+        lists_ret[j].append(
+            tensor_conversion_registry.convert(lists[j][i], dtype=dtype)
+        )
+    else:
+      # Convert everything to the found dtype.
+      for j in range(len(lists)):
+        lists_ret[j].append(
+            tensor_conversion_registry.convert(lists[j][i], dtype=dtype)
+        )
+    types.append(dtype.as_datatype_enum)
+  return types, lists_ret
+
+
+def _is_keras_symbolic_tensor(x):
+  return hasattr(x, "graph") and getattr(x.graph, "name", None) == "keras_graph"

tensorflow/python/eager/executor.py

@@ -0,0 +1,77 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Executor for eager execution."""
+
+from tensorflow.python import pywrap_tfe
+
+
+class Executor(object):
+  """A class for handling eager execution.
+
+  The default behavior for asynchronous execution is to serialize all ops on
+  a single thread. Having different `Executor` objects in different threads
+  enables executing ops asynchronously in parallel:
+
+  ```python
+  def thread_function():
+    executor = executor.Executor(enable_async=True):
+    context.set_executor(executor)
+
+  a = threading.Thread(target=thread_function)
+  a.start()
+  b = threading.Thread(target=thread_function)
+  b.start()
+  ```
+  """
+
+  __slots__ = ["_handle"]
+
+  def __init__(self, handle):
+    self._handle = handle
+
+  def __del__(self):
+    try:
+      self.wait()
+      pywrap_tfe.TFE_DeleteExecutor(self._handle)
+    except TypeError:
+      # Suppress some exceptions, mainly for the case when we're running on
+      # module deletion. Things that can go wrong include the pywrap module
+      # already being unloaded, self._handle. no longer being
+      # valid, and so on. Printing warnings in these cases is silly
+      # (exceptions raised from __del__ are printed as warnings to stderr).
+      pass  # 'NoneType' object is not callable when the handle has been
+      # partially unloaded.
+
+  def is_async(self):
+    return pywrap_tfe.TFE_ExecutorIsAsync(self._handle)
+
+  def handle(self):
+    return self._handle
+
+  def wait(self):
+    """Waits for ops dispatched in this executor to finish."""
+    pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)
+
+  def clear_error(self):
+    """Clears errors raised in this executor during execution."""
+    pywrap_tfe.TFE_ExecutorClearError(self._handle)
+
+
+def new_executor(enable_async,
+                 enable_streaming_enqueue=True,
+                 in_flight_nodes_limit=0):
+  handle = pywrap_tfe.TFE_NewExecutor(enable_async, enable_streaming_enqueue,
+                                      in_flight_nodes_limit)
+  return Executor(handle)

tensorflow/python/eager/forwardprop.py

@@ -0,0 +1,487 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Utilities for forward-mode automatic differentiation."""
+
+import functools
+import threading
+
+from tensorflow.core.function.polymorphism import function_cache
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import backprop_util
+from tensorflow.python.eager import execute
+from tensorflow.python.eager import forwardprop_util
+from tensorflow.python.eager.polymorphic_function import tracing_compilation
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.parallel_for import control_flow_ops
+from tensorflow.python.ops.unconnected_gradients import UnconnectedGradients
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+# Dictionary mapping from op names to special-cased jvp functions. Otherwise
+# backward functions are transposed on the tape.
+_SPECIAL_CASES = {}
+
+
+def _identity_jvp(attr_tuple, inputs, outputs, tangents):
+  # Special-cased mostly for resource handles, where creating ones Tensors from
+  # handle data for transposing the backward function on the tape is error-prone
+  # (even if we get good handle data, partially defined shapes are an issue).
+  del attr_tuple, inputs, outputs
+  return [array_ops.identity(t) for t in tangents]
+
+
+_SPECIAL_CASES["Identity"] = _identity_jvp
+
+
+def _read_variable_jvp(attr_tuple, inputs, outputs, tangents):
+  # Like for Identity, this special case means we don't need to create
+  # variable-shaped Tensors from resource handles.
+  del attr_tuple, inputs, outputs
+  return [array_ops.identity(t) for t in tangents]
+
+
+_SPECIAL_CASES["ReadVariableOp"] = _read_variable_jvp
+
+
+_TRACE_COUNT_CONSISTENCY_LOCK = threading.Lock()
+# Map from op names to number of traces of _jvp_helper. Used to cap the number
+# of traces due to shape differences while still specializing where possible.
+_TRACE_COUNT = {}
+
+
+def _jvp_helper(op_name, attr_tuple, inputs, outputs, tangents):
+  """Computes a Jacobian-vector product for an op.
+
+  Note that this function would be wasteful if executed eagerly. It runs the
+  backward gradient function and throws away the result just to record its
+  operations on a GradientTape. These unused ops are pruned away when this
+  function is traced.
+
+  Args:
+    op_name: A string, the type of operation being executed.
+    attr_tuple: Attributes of the operation.
+    inputs: A flat list of input Tensors to the operation.
+    outputs: A flat list of output Tensors from the operation.
+    tangents: A flat list of Tensors, same shape as `inputs`.
+
+  Returns:
+    A flat list of tangents corresponding to `outputs`.
+  """
+  with _TRACE_COUNT_CONSISTENCY_LOCK:
+    # Just make sure writes don't clobber each other's increments; reads in
+    # _jvp_dispatch do not lock.
+    _TRACE_COUNT[op_name] = _TRACE_COUNT.get(op_name, 0) + 1
+
+  special_case = _SPECIAL_CASES.get(op_name, None)
+  if special_case is not None:
+    return special_case(attr_tuple, inputs, outputs, tangents)
+  if not outputs:
+    # tape.gradients([], inputs) doesn't make much sense
+    return []
+  # Generally inner GradientTapes won't function while outer accumulators are
+  # recording. We temporarily reset forwardprop state to allow GradientTapes to
+  # function here.
+  with forwardprop_util.push_forwardprop_state():
+    trainable_inputs = []
+    trainable_indices = []
+    nontrivial_tangents = []
+    for input_index, tensor in enumerate(inputs):
+      if backprop_util.IsTrainable(tensor):
+        trainable_inputs.append(tensor)
+        trainable_indices.append(input_index)
+        nontrivial_tangents.append(tangents[input_index])
+
+    with backprop.GradientTape() as transpose_tape:
+      with backprop.GradientTape() as backfunc_tape:
+        backfunc_tape.watch(trainable_inputs)
+        execute.record_gradient(op_name, inputs, attr_tuple, outputs)
+
+      forwardprop_aids = []
+      trainable_outputs = []
+      nontrivial_output_indices = []
+      for output_index, output in enumerate(outputs):
+        if backprop_util.IsTrainable(output):
+          forwardprop_aids.append(
+              array_ops.ones_like(output, name="unused_forwardprop_aid"))
+          trainable_outputs.append(output)
+          nontrivial_output_indices.append(output_index)
+
+      transpose_tape.watch(forwardprop_aids)
+      grads = backfunc_tape.gradient(
+          trainable_outputs,
+          trainable_inputs,
+          forwardprop_aids,
+          unconnected_gradients=UnconnectedGradients.ZERO)
+    nontrivial_output_tangents = transpose_tape.gradient(
+        grads, forwardprop_aids, output_gradients=nontrivial_tangents)
+    output_tangents = [None] * len(outputs)
+    for index, tangent in zip(nontrivial_output_indices,
+                              nontrivial_output_tangents):
+      output_tangents[index] = tangent
+    return output_tangents
+
+
+def _jvp_helper_wrapper(op_name, attr_tuple, inputs, outputs, tangents,
+                        use_batch):
+  """Computes a batch of Jacobian-vector product for an op.
+
+  Args:
+    op_name: A string, the type of operation being executed.
+    attr_tuple: Attributes of the operation.
+    inputs: A flat list of input Tensors to the operation.
+    outputs: A flat list of output Tensors from the operation.
+    tangents: A flat list of Tensors, compatible with shape `[None] +
+      input_shape`.
+    use_batch: A bool, True to vetorize over batch of tangents of shape `[None]
+      + input_shape`.
+
+  Returns:
+    A flat list of tangents compatible with `outputs`
+    or `[None] + output_shape`.
+
+  Raises:
+    ValueError: if tangent shapes are not compatible with input shapes.
+  """
+  if use_batch:
+    for primal, tangent in zip(inputs, tangents):
+      if not tangent.shape.is_compatible_with([None] + primal.shape):
+        raise ValueError("Tangent {} was expected to be of shape "
+                         "{} but is instead of shape {}".format(
+                             tangent, [None] + primal.shape, tangent.shape))
+
+    return control_flow_ops.vectorized_map(
+        functools.partial(_jvp_helper, op_name, attr_tuple, inputs, outputs),
+        tangents,
+    )
+  return _jvp_helper(op_name, attr_tuple, inputs, outputs, tangents)
+
+
+# TODO(allenl): reduce_retracing for gradients which rely on static
+# shape information are underspecialized. We may want hand-written forward
+# implementations, or a more satisfying story about how we re-specialize
+# gradients which were traced with relaxed shapes (e.g. use conds instead of
+# trace-time Python logic).
+#
+# Using function.defun rather than def_function.function avoids
+# tf.config.run_functions_eagerly(True). `_jvp_helper` doesn't successfully run
+# eagerly (infinite recursion), and even if it did it would use extra memory and
+# run unnecessary computation. The function does not create variables, so the
+# two symbols are otherwise equivalent.
+_jvp_function_cache = function_cache.FunctionCache()
+_jvp_relaxed_config = tracing_compilation.TracingOptions(
+    _jvp_helper_wrapper,
+    name="_jvp_relaxed_shapes",
+    reduce_retracing=True,
+    function_cache=_jvp_function_cache,
+)
+
+_jvp_exact_config = tracing_compilation.TracingOptions(
+    _jvp_helper_wrapper,
+    name="_jvp_exact_shapes",
+    reduce_retracing=False,
+    function_cache=_jvp_function_cache,
+)
+
+# The maximum number of exact-shape traces to perform for a single op before
+# switching to shape relaxation.
+_TRACE_COUNT_LIMIT = 32
+
+
+def _jvp_dispatch(op_name,
+                  attr_tuple,
+                  inputs,
+                  outputs,
+                  tangents,
+                  use_batch=False):
+  """Determine which forwardprop function to call."""
+  # Note that this _TRACE_COUNT read races with writes. That's fine, it just
+  # means we may trace a few more exact shapes before moving on to relaxation.
+  if _TRACE_COUNT.get(op_name, 0) < _TRACE_COUNT_LIMIT:
+    config = _jvp_exact_config
+  else:
+    config = _jvp_relaxed_config
+  return tracing_compilation.call_function(
+      (op_name, attr_tuple, inputs, outputs, tangents, use_batch),
+      tracing_options=config,
+  )
+
+
+pywrap_tfe.TFE_Py_RegisterJVPFunction(_jvp_dispatch)
+
+
+@tf_export("autodiff.ForwardAccumulator", v1=[])
+class ForwardAccumulator():
+  """Computes Jacobian-vector products ("JVP"s) using forward-mode autodiff.
+
+  Compare to `tf.GradientTape` which computes vector-Jacobian products ("VJP"s)
+  using reverse-mode autodiff (backprop). Reverse mode is more attractive when
+  computing gradients of a scalar-valued function with respect to many inputs
+  (e.g. a neural network with many parameters and a scalar loss). Forward mode
+  works best on functions with many outputs and few inputs. Since it does not
+  hold on to intermediate activations, it is much more memory efficient than
+  backprop where it is applicable.
+
+  Consider a simple linear regression:
+
+  >>> x = tf.constant([[2.0, 3.0], [1.0, 4.0]])
+  >>> targets = tf.constant([[1.], [-1.]])
+  >>> dense = tf.keras.layers.Dense(1)
+  >>> dense.build([None, 2])
+  >>> with tf.autodiff.ForwardAccumulator(
+  ...    primals=dense.kernel,
+  ...    tangents=tf.constant([[1.], [0.]])) as acc:
+  ...   loss = tf.reduce_sum((dense(x) - targets) ** 2.)
+  >>> acc.jvp(loss)
+  <tf.Tensor: shape=(), dtype=float32, numpy=...>
+
+  The example has two variables containing parameters, `dense.kernel` (2
+  parameters) and `dense.bias` (1 parameter). Considering the training data `x`
+  as a constant, this means the Jacobian matrix for the function mapping from
+  parameters to loss has one row and three columns.
+
+  With forwardprop, we specify a length-three vector in advance which multiplies
+  the Jacobian. The `primals` constructor argument is the parameter (a
+  `tf.Tensor` or `tf.Variable`) we're specifying a vector for, and the
+  `tangents` argument is the "vector" in Jacobian-vector product. If our goal is
+  to compute the entire Jacobian matrix, forwardprop computes one column at a
+  time while backprop computes one row at a time. Since the Jacobian in the
+  linear regression example has only one row, backprop requires fewer
+  invocations:
+
+  >>> x = tf.constant([[2.0, 3.0], [1.0, 4.0]])
+  >>> targets = tf.constant([[1.], [-1.]])
+  >>> dense = tf.keras.layers.Dense(1)
+  >>> dense.build([None, 2])
+  >>> loss_fn = lambda: tf.reduce_sum((dense(x) - targets) ** 2.)
+  >>> kernel_fprop = []
+  >>> with tf.autodiff.ForwardAccumulator(
+  ...     dense.kernel, tf.constant([[1.], [0.]])) as acc:
+  ...   kernel_fprop.append(acc.jvp(loss_fn()))
+  >>> with tf.autodiff.ForwardAccumulator(
+  ...     dense.kernel, tf.constant([[0.], [1.]])) as acc:
+  ...   kernel_fprop.append(acc.jvp(loss_fn()))
+  >>> with tf.autodiff.ForwardAccumulator(dense.bias, tf.constant([1.])) as acc:
+  ...   bias_fprop = acc.jvp(loss_fn())
+  >>> with tf.GradientTape() as tape:
+  ...   loss = loss_fn()
+  >>> kernel_grad, bias_grad = tape.gradient(loss, (dense.kernel, dense.bias))
+  >>> np.testing.assert_allclose(
+  ...     kernel_grad, tf.stack(kernel_fprop)[:, tf.newaxis])
+  >>> np.testing.assert_allclose(bias_grad, bias_fprop[tf.newaxis])
+
+  Implicit in the `tape.gradient` call is a length-one vector which
+  left-multiplies the Jacobian, a vector-Jacobian product.
+
+  `ForwardAccumulator` maintains JVPs corresponding primal tensors it is
+  watching, derived from the original `primals` specified in the constructor. As
+  soon as a primal tensor is deleted, `ForwardAccumulator` deletes the
+  corresponding JVP.
+
+  `acc.jvp(x)` retrieves `acc`'s JVP corresponding to the primal tensor `x`. It
+  does not perform any computation. `acc.jvp` calls can be repeated as long as
+  `acc` is accessible, whether the context manager is active or not. New JVPs
+  are only computed while the context manager is active.
+
+  Note that `ForwardAccumulator`s are always applied in the order their context
+  managers were entered, so inner accumulators will not see JVP computation from
+  outer accumulators. Take higher-order JVPs from outer accumulators:
+
+  >>> primal = tf.constant(1.1)
+  >>> with tf.autodiff.ForwardAccumulator(primal, tf.constant(1.)) as outer:
+  ...   with tf.autodiff.ForwardAccumulator(primal, tf.constant(1.)) as inner:
+  ...     primal_out = primal ** tf.constant(3.5)
+  >>> inner_jvp = inner.jvp(primal_out)
+  >>> inner_jvp  # 3.5 * 1.1 ** 2.5
+  <tf.Tensor: shape=(), dtype=float32, numpy=4.4417057>
+  >>> outer.jvp(inner_jvp)  # 3.5 * 2.5 * 1.1 ** 1.5
+  <tf.Tensor: shape=(), dtype=float32, numpy=10.094786>
+
+  Reversing the collection in the last line to instead retrieve
+  `inner.jvp(outer.jvp(primal_out))` will not work.
+
+  Strict nesting also applies to combinations of `ForwardAccumulator` and
+  `tf.GradientTape`. More deeply nested `GradientTape` objects will ignore the
+  products of outer `ForwardAccumulator` objects. This allows (for example)
+  memory-efficient forward-over-backward computation of Hessian-vector products,
+  where the inner `GradientTape` would otherwise hold on to all intermediate
+  JVPs:
+
+  >>> v = tf.Variable([1., 2.])
+  >>> with tf.autodiff.ForwardAccumulator(
+  ...     v,
+  ...     # The "vector" in Hessian-vector product.
+  ...     tf.constant([1., 0.])) as acc:
+  ...   with tf.GradientTape() as tape:
+  ...     y = tf.reduce_sum(v ** 3.)
+  ...   backward = tape.gradient(y, v)
+  >>> backward  # gradient from backprop
+  <tf.Tensor: shape=(2,), dtype=float32, numpy=array([ 3., 12.], dtype=float32)>
+  >>> acc.jvp(backward)  # forward-over-backward Hessian-vector product
+  <tf.Tensor: shape=(2,), dtype=float32, numpy=array([6., 0.], dtype=float32)>
+  """
+
+  def __init__(self, primals, tangents):
+    """Specify tensors to watch and their Jacobian-vector products.
+
+    Mathematically, `tangents` is a vector right-multiplying the Jacobian matrix
+    (a Jacobian-vector product) for the function computed while this accumulator
+    is active. Since JVPs are computed in forward mode as the computation
+    happens, this vector must be supplied in advance.
+
+    Listing a single tensor multiple times in `primals` raises an
+    exception. Excluding a tensor from `primals` is equivalent to watching it
+    with a tangent tensor of zeros.
+
+    Args:
+      primals: A tensor or nested structure of tensors to watch.
+      tangents: A tensor or nested structure of tensors, with the same nesting
+        structure as `primals`, with each element being a vector with the same
+        size as the corresponding primal element.
+
+    Raises:
+      ValueError: If the same tensor or variable is specified multiple times in
+        `primals`.
+    """
+    self._accumulator = pywrap_tfe.TFE_Py_ForwardAccumulatorNew(False)
+    self._recording = False
+    primal_ids = set()
+    for primal in nest.flatten(primals):
+      if id(primal) in primal_ids:
+        raise ValueError(
+            "Tensor {} was specified as a primal multiple times. This may "
+            "indicate an error. If it was intended, please sum the "
+            "corresponding tangents.")
+      primal_ids.add(id(primal))
+    self._watch(primals, tangents)
+
+  def __enter__(self):
+    self._push_accumulator()
+    return self
+
+  def __exit__(self, typ, value, traceback):
+    if self._recording:
+      self._pop_accumulator()
+
+  def _push_accumulator(self):
+    if self._recording:
+      raise ValueError("Accumulator is already recording.")
+    pywrap_tfe.TFE_Py_ForwardAccumulatorSetAdd(self._accumulator)
+    self._recording = True
+
+  def _pop_accumulator(self):
+    if not self._recording:
+      raise ValueError("Accumulator is not recording.")
+    pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(self._accumulator)
+    self._recording = False
+
+  def _watch(self, primals, tangents):
+    """Ensures that `primals` are being traced by this accumulator.
+
+    Mathematically, `tangents` is a vector right-multiplying the Jacobian matrix
+    (a Jacobian-vector product) for the function computed while this accumulator
+    is active. Since JVPs are computed in forward mode as the computation
+    happens, this vector must be supplied in advance.
+
+    Watching a single tensor multiple times sums each of its `tangents`. Any
+    un-watched tensor has zeros for its tangent vector.
+
+    Args:
+      primals: A Tensor or list of Tensors.
+      tangents: A Tensor or list of Tensors matching `primals`.
+    """
+
+    def _watch(primal, tangent):
+      if not primal.dtype.is_floating:
+        logging.log_first_n(
+            logging.WARN, "The dtype of the watched primal must be "
+            "floating (e.g. tf.float32), got %r", 5, primal.dtype)
+      tangent = ops.convert_to_tensor(tangent, dtype=primal.dtype)
+      if hasattr(primal, "handle"):
+        # Run convert_to_tensor to get the captured handle from whichever
+        # function we're running if necessary.
+        primal = ops.convert_to_tensor(primal.handle)
+      pywrap_tfe.TFE_Py_ForwardAccumulatorWatch(self._accumulator, primal,
+                                                tangent)
+
+    nest.map_structure(_watch, primals, tangents)
+
+  def jvp(self, primals, unconnected_gradients=UnconnectedGradients.NONE):
+    """Fetches the Jacobian-vector product computed for `primals`.
+
+    Note that this method performs no computation, and simply looks up a JVP
+    that was already computed (unlike backprop using a `tf.GradientTape`, where
+    the computation happens on the call to `tape.gradient`).
+
+    Args:
+      primals: A watched Tensor or structure of Tensors to fetch the JVPs for.
+      unconnected_gradients: A value which can either hold 'none' or 'zero' and
+        alters the value which will be returned if no JVP was computed for
+        `primals`. The possible values and effects are detailed in
+        'tf.UnconnectedGradients' and it defaults to 'none'.
+
+    Returns:
+      Tensors with the same shapes and dtypes as `primals`, or None if no JVP
+      is available.
+    """
+    unconnected_gradients = UnconnectedGradients(unconnected_gradients)
+    if self._accumulator is None:
+      raise ValueError("Called jvp() without first tracing anything.")
+
+    def _fetch_jvp(tensor):
+      if hasattr(tensor, "handle"):
+        unwrapped_tensor = ops.convert_to_tensor(tensor.handle)
+      else:
+        unwrapped_tensor = tensor
+      result = pywrap_tfe.TFE_Py_ForwardAccumulatorJVP(self._accumulator,
+                                                       unwrapped_tensor)
+      if result is None and unconnected_gradients == UnconnectedGradients.ZERO:
+        result = array_ops.zeros_like(tensor)
+      return result
+
+    return nest.map_structure(_fetch_jvp, primals)
+
+  @classmethod
+  def _batch_accumulator(cls, primals, tangents):
+    """Factory constructor to test accumulator on batches of tangents.
+
+    Args:
+      primals: A tensor or nested structure of tensors to watch.
+      tangents: A tensor or nested structure of tensors, with the same nesting
+        structure as `primals`, with each element being a vector with compatible
+        shape `[None] + primal.shape` of the corresponding primal element.
+
+    Returns:
+      A batch accumulator object.
+    """
+    acc = super(ForwardAccumulator, cls).__new__(cls, primals, tangents)
+    acc._recording = False
+    acc._accumulator = pywrap_tfe.TFE_Py_ForwardAccumulatorNew(True)
+    primal_ids = set()
+    for primal, tangent in zip(nest.flatten(primals), nest.flatten(tangents)):
+      tangent.shape.assert_is_compatible_with(
+          tensor_shape.TensorShape([None]) + primal.shape)
+      if id(primal) in primal_ids:
+        raise ValueError(
+            "Tensor {} was specified as a primal multiple times. This may "
+            "indicate an error. If it was intended, please sum the "
+            "corresponding tangents.")
+      primal_ids.add(id(primal))
+    acc._watch(primals, tangents)
+    return acc

tensorflow/python/eager/forwardprop_util.py

@@ -0,0 +1,74 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Utilities for managing forward accumulators.
+
+A separate file from forwardprop.py so that functions can use these utilities.
+"""
+
+import collections
+import contextlib
+
+from tensorflow.python import pywrap_tfe
+
+
+class TangentInfo(
+    collections.namedtuple("TangentInfo", ["indices", "tangents"])):
+  """Packed forward accumulator state. The return value of `pack_tangents`."""
+
+  def __new__(cls, indices=None, tangents=None):
+    if indices is None:
+      indices = ()
+    if tangents is None:
+      tangents = []
+    return super(TangentInfo, cls).__new__(cls, indices, tangents)
+
+
+def pack_tangents(tensors):
+  """Packs forward accumulator state into a TangentInfo tuple.
+
+  Args:
+    tensors: A flat list of Tensors to pack forward accumulator state for.
+
+  Returns:
+    A tuple of (indices, tangents):
+      indices: A sequence of sequences of two-element tuples. Each forward
+        accumulator is represented as a sequence of tuples with (primal_index,
+        jvp_index). Both integers index into the concatenated `tensors + jvps`
+        array.
+      tangents: A flat list of Tensors. Best interpreted as a sequence to be
+        appended to `tensors`.
+  """
+  return TangentInfo(*pywrap_tfe.TFE_Py_PackJVPs(tensors))
+
+
+@contextlib.contextmanager
+def push_forwardprop_state():
+  """Temporarily push or pop transient state for accumulators in the active set.
+
+  Allows an accumulator which is currently processing an operation to
+  temporarily reset its state. This is useful when building forwardprop versions
+  of functions, where an accumulator will trigger function building and then
+  must process captured symbolic tensors while building it. Without pushing and
+  popping, accumulators ignore operations executed as a direct result of their
+  own jvp computations.
+
+  Yields:
+    None (used for its side effect).
+  """
+  try:
+    pywrap_tfe.TFE_Py_ForwardAccumulatorPushState()
+    yield
+  finally:
+    pywrap_tfe.TFE_Py_ForwardAccumulatorPopState()

tensorflow/python/eager/function.py

@@ -0,0 +1,37 @@
+# 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.
+# ==============================================================================
+"""Supports old symbols supplied by this file while the code is refactored."""
+
+# pylint:disable=unused-import,g-bad-import-order
+
+# TODO(b/243822285): Reduce this list as much as possible.
+# Constants
+from tensorflow.python.eager.polymorphic_function.concrete_function import _BACKWARD_PREFIX
+from tensorflow.python.eager.polymorphic_function.concrete_function import _FORWARD_PREFIX
+from tensorflow.python.eager.polymorphic_function.concrete_function import _INFERENCE_PREFIX
+
+# Function Classes
+from tensorflow.python.eager.polymorphic_function.concrete_function import ConcreteFunction
+from tensorflow.python.eager.polymorphic_function.atomic_function import from_func_graph
+from tensorflow.python.eager.polymorphic_function.atomic_function import AtomicFunction
+
+# Utilities
+from tensorflow.python.eager.polymorphic_function.tf_method_target import TfMethodTarget
+from tensorflow.python.eager.polymorphic_function.concrete_function import _inference_name
+
+# TODO(b/244360504): Remove in favor of graph transformation API.
+# QUARANTINED - Function Callback Modification API
+from tensorflow.python.eager.polymorphic_function.transform import FUNC_GRAPH_TRANSFORMS
+from tensorflow.python.eager.polymorphic_function.transform import CONCRETE_FUNCTION_CALLBACKS

tensorflow/python/eager/gen_gradient_input_output_exclusions.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.
+r"""Script to generate inputs/outputs exclusion lists for GradientTape.
+
+To use this script:
+
+bazel run tensorflow/python/eager:gen_gradient_input_output_exclusions -- \
+  $PWD/tensorflow/python/eager/pywrap_gradient_exclusions.cc
+"""
+
+import argparse
+
+from tensorflow.python.eager import gradient_input_output_exclusions
+
+
+def main(output_file):
+  with open(output_file, "w") as fp:
+    fp.write(gradient_input_output_exclusions.get_contents())
+
+
+if __name__ == "__main__":
+  arg_parser = argparse.ArgumentParser()
+  arg_parser.add_argument("output", metavar="O", type=str, help="Output file.")
+  args = arg_parser.parse_args()
+  main(args.output)

tensorflow/python/eager/gradient_input_output_exclusions.py

@@ -0,0 +1,367 @@
+# 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.
+"""Code to generate inputs/outputs exclusion lists for GradientTape."""
+
+import sys
+
+import gast
+
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import cfg
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import transformer
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis import liveness
+from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+
+_GENERATED_FILE_HEADER = """/* 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.
+==============================================================================*/
+
+// Inputs/Outputs exclusion lists for GradientTape.
+//
+// This file is MACHINE GENERATED! Do not edit.
+// Generated by: tensorflow/python/eager/gen_gradient_input_output_exclusions.py
+"""
+
+_INCLUDES = """
+#include "tensorflow/python/eager/pywrap_gradient_exclusions.h"
+
+#include "absl/types/optional.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/gtl/flatset.h"
+
+using tensorflow::string;
+
+namespace {
+// Keep static data in a format that's easy to init statically.
+struct OpIndexInfo {
+  const char *op_name;
+  int num_indices;
+  std::array<int, 4> unused_indices;
+};
+
+// Helper function to initialize FlatMap<string,FlatSet> from OpIndexInfo.
+template <typename T>
+auto OpGradientInfoInit(const T &a) {
+  auto *m = new tensorflow::gtl::FlatMap<string, tensorflow::gtl::FlatSet<int>>;
+  for (const auto &item : a) {
+    m->emplace(string(item.op_name),
+               tensorflow::gtl::FlatSet<int>(
+                   item.unused_indices.begin(),
+                   item.unused_indices.begin() + item.num_indices));
+  }
+  return m;
+}
+}  // namespace
+"""
+
+_EXCLUDED_OPS = [
+    # Composite ops with custom gradient functions.
+    "If",
+    "StatelessIf",
+    "While",
+    "StatelessWhile",
+    "Case",
+    # TF Lite. These ops only appear in OSS.
+    # TODO(srbs): Find a better way to filter these out.
+    "AudioMicrofrontend",
+    # DTensor Ops with custom gradient functions.
+    # Note that these ops only appear in OSS, and fails the test in OSS.
+    "CopyToMesh",
+    "CopyToMeshGrad",
+    "Relayout",
+    "RelayoutLike",
+    # Debug ops that similarly only appear in OSS, and fails the test in OSS.
+    "DebugGradientIdentity",
+    "DebugGradientRefIdentity",
+    "DebugIdentityV2",
+]
+
+
+class _SubscriptUseTracker(transformer.Base):
+  """Track uses of composite names, excluding certain names when subscripted."""
+
+  def __init__(self, ctx, exclude_when_subscripted):
+    super(_SubscriptUseTracker, self).__init__(ctx)
+    self.exclude = exclude_when_subscripted
+    self.reads = set()
+    self.complex_reads = set()
+
+  def visit_Attribute(self, node):
+    """Visits attribute nodes in the AST."""
+    if anno.hasanno(node, anno.Basic.QN):
+      qn = anno.getanno(node, anno.Basic.QN)
+      if isinstance(node.ctx, gast.Load):
+        self.reads.add(qn)
+    node = self.generic_visit(node)
+    return node
+
+  def visit_Subscript(self, node):
+    """Visits nodes with subscript in the AST."""
+    s = node.slice
+    if anno.hasanno(node, anno.Basic.QN):
+      qn = anno.getanno(node, anno.Basic.QN)
+      if isinstance(node.ctx, gast.Load):
+        self.reads.add(qn)
+    elif isinstance(s, (gast.Tuple, gast.Slice)):
+      if anno.hasanno(node.value, anno.Basic.QN):
+        self.complex_reads.add(anno.getanno(node.value, anno.Basic.QN))
+    value_qn = anno.getanno(node.value, anno.Basic.QN, None)
+    if value_qn in self.exclude:
+      node.value = self.generic_visit(node.value)
+    else:
+      node.value = self.visit(node.value)
+    node.slice = self.visit(s)
+    return node
+
+
+class _FunctionCallsTracker(transformer.Base):
+  """Tracks any function calls made with a given first argument name."""
+
+  def __init__(self, ctx, first_argument_name):
+    super(_FunctionCallsTracker, self).__init__(ctx)
+    self.first_argument_name = first_argument_name
+    self.calls = set()
+
+  def visit_Name(self, node):
+    node = self.generic_visit(node)
+    if isinstance(node.ctx, gast.Load) and node.id in self.ctx.info.namespace:
+      anno.setanno(node, "static_value", self.ctx.info.namespace[node.id])
+    return node
+
+  def visit_Attribute(self, node):
+    node = self.generic_visit(node)
+    parent_val = anno.getanno(node.value, "static_value", default=None)
+    if parent_val is not None:
+      if hasattr(parent_val, node.attr):
+        anno.setanno(node, "static_value", getattr(parent_val, node.attr))
+    return node
+
+  def visit_Call(self, node):
+    node = self.generic_visit(node)
+    if (node.args and anno.getanno(node.args[0], anno.Basic.QN,
+                                   None) == self.first_argument_name):
+      fn_object = anno.getanno(node.func, "static_value", None)
+      if fn_object is not None:
+        self.calls.add(fn_object)
+    return node
+
+
+_ALL = object()
+
+
+def _live_tensors(f, attr_name="inputs"):
+  """Returns the indices of the used inputs.
+
+  Note: This currently only handles direct index accesses e.g. op.inputs[1].
+  If the function has slicing or list comprehension on attr_name then returns
+  _ALL. This ensure that this is correct even if inefficient.
+
+  Args:
+    f: A grad function, taking the op as first argument.
+    attr_name: op attr to track. "inputs" or "outputs".
+
+  Returns:
+    Either one of:
+      * set of integers representing individual indices of inputs used
+      * the value _ALL, if indices are used but cannot be determined which
+      * empty set, if no inputs are used
+  """
+  node, _ = parser.parse_entity(f, ())
+  entity_info = transformer.EntityInfo(
+      name=f.__name__,
+      source_code=None,
+      source_file=None,
+      future_features=(),
+      namespace=sys.modules[f.__module__].__dict__)
+  ctx = transformer.Context(entity_info, None, None)
+
+  graphs = cfg.build(node)
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+  node = reaching_fndefs.resolve(node, ctx, graphs)
+  node = liveness.resolve(node, ctx, graphs)
+
+  op_arg_name = anno.getanno(node.args.args[0], anno.Basic.QN)
+  op_inputs_outputs_name = qual_names.QN(op_arg_name, attr=attr_name)
+
+  special_tracker = _SubscriptUseTracker(ctx, (op_inputs_outputs_name,))
+  node = special_tracker.visit(node)
+
+  live_vars_in = anno.getanno(node.body[0], anno.Static.LIVE_VARS_IN)
+  inputs_outputs_used_qns = set()
+  for v in special_tracker.complex_reads:
+    # Complicated patterns like op.inputs[:3]. Could be smarter about them
+    # if they matter much.
+    if v == op_inputs_outputs_name:
+      return _ALL
+  for v in live_vars_in:
+    if v in special_tracker.reads:
+      if (v.has_subscript() and v.parent == op_inputs_outputs_name):
+        inputs_outputs_used_qns.add(v)
+      elif v == op_inputs_outputs_name:
+        # When op.{attr_name} is used directly, assume all tensors are
+        # used for now. In that case, no point digging further.
+        # TODO(mdan): We can descend into tuple expansions.
+        return _ALL
+
+  function_calls_tracker = _FunctionCallsTracker(ctx, op_arg_name)
+  node = function_calls_tracker.visit(node)
+
+  input_output_indices = set()
+
+  for called_f in function_calls_tracker.calls:
+    child_indices = _live_tensors(called_f, attr_name=attr_name)
+    if child_indices is _ALL:
+      return _ALL
+    input_output_indices |= child_indices
+
+  for v in inputs_outputs_used_qns:
+    assert v.has_subscript()
+    _, subscript = v.qn
+    if not subscript.is_simple():
+      # Not a number, assuming it can be anything.
+      return _ALL
+    subscript_val, = subscript.qn
+    if (not isinstance(subscript_val, qual_names.Literal) and
+        not isinstance(subscript_val.value, int)):
+      # Not a number, assuming it can be anything.
+      return _ALL
+    input_output_indices.add(subscript_val.value)
+  return input_output_indices
+
+
+def _get_num_inputs_outputs(op_type):
+  """Returns (num_inputs, num_outputs).
+
+  Args:
+    op_type: String. The type of the Operation. Used to lookup the op in the
+      registry.
+
+  Returns:
+    (num_inputs, num_outputs), for either num_inputs or num_outputs if the value
+    can't be statically inferred from the OpDef alone or of the OpDef lookup
+    fails, -1 is returned.
+  """
+
+  def _is_list_arg(arg):
+    return arg.number_attr or arg.type_list_attr
+
+  def _count_args(arg_defs):
+    for arg in arg_defs:
+      if _is_list_arg(arg):
+        # Op has list type args which could be variable.
+        return -1
+    return len(arg_defs)
+
+  op_def = op_def_registry.get(op_type)
+  if not op_def:
+    return -1, -1
+  return _count_args(op_def.input_arg), _count_args(op_def.output_arg)
+
+
+def get_entries(attr_name):
+  """Returns the dict of entries.
+
+  Each entry is of the form {op_name, {true|false, indices}}
+
+  true: All values are unused.
+  false: `indices` are the only unused indices.
+
+  Note: ops for which all values are used are not printed.
+
+  Args:
+    attr_name: inputs or outputs.
+
+  Returns:
+    A dict from op_type to formatted entry in the dict.
+  """
+  assert attr_name in ["inputs", "outputs"]
+  entries = {}
+  for op_type in ops._gradient_registry.list():  # pylint: disable=protected-access
+    if op_type in _EXCLUDED_OPS:
+      continue
+    num_values = _get_num_inputs_outputs(op_type)[0 if attr_name ==
+                                                  "inputs" else 1]
+    gradient_fn = ops._gradient_registry.lookup(op_type)  # pylint: disable=protected-access
+    if gradient_fn is None:
+      # NotDifferentiable
+      if num_values != -1:
+        entries[op_type] = "{\"%s\"}," % op_type
+      continue
+    used_tensors = _live_tensors(gradient_fn, attr_name=attr_name)
+    if used_tensors is _ALL:
+      continue
+    elif not used_tensors:
+      entries[op_type] = "{\"%s\"}," % op_type
+    else:
+      all_tensors = set(range(num_values))
+      unused_tensors = all_tensors - used_tensors
+      if unused_tensors:
+        unused_tensor_list = sorted(list(unused_tensors))
+        entries[op_type] = "{\"%s\", %d, {%s}}," % (
+            op_type, len(unused_tensor_list), ", ".join(
+                str(i) for i in unused_tensor_list))
+  return entries
+
+
+def get_function(name, entries):
+  """Generates lookup function with given name and lookup table entries."""
+  contents = """
+absl::optional<tensorflow::gtl::FlatSet<int>> {name}(
+    const tensorflow::string &op_name) {{
+  static std::array<OpIndexInfo, {count}> a = {{{{
+""".format(
+    name=name, count=len(entries) + 1)
+  contents += "      "
+  contents += "\n      ".join(entries[op_type] for op_type in sorted(entries))
+  contents += "\n      {\"VarHandleOp\"},"
+  contents += """
+  }};
+  static const auto &m = *OpGradientInfoInit(a);
+
+  auto it = m.find(op_name);
+  if (it != m.end()) {
+    return it->second;
+  }
+  return absl::nullopt;
+}
+"""
+  return contents
+
+
+def get_contents():
+  """Returns contents for the generated file."""
+  contents = ""
+  contents += _GENERATED_FILE_HEADER + _INCLUDES
+  contents += get_function("OpGradientUnusedInputIndices",
+                           get_entries("inputs"))
+  contents += get_function("OpGradientUnusedOutputIndices",
+                           get_entries("outputs"))
+  return contents

tensorflow/python/eager/gradient_input_output_exclusions_test.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+"""Ensures that pywrap_gradient_exclusions.cc is up-to-date."""
+
+import os
+
+from tensorflow.python.eager import gradient_input_output_exclusions
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+
+
+class GradientInputOutputExclusionsTest(test.TestCase):
+
+  def testGeneratedFileMatchesHead(self):
+    expected_contents = gradient_input_output_exclusions.get_contents()
+    filename = os.path.join(
+        resource_loader.get_root_dir_with_all_resources(),
+        resource_loader.get_path_to_datafile("pywrap_gradient_exclusions.cc"))
+    actual_contents = file_io.read_file_to_string(filename)
+
+    # On windows, one or both of these strings may have CRLF line endings.
+    # To make sure, sanitize both:
+    sanitized_actual_contents = actual_contents.replace("\r", "")
+    sanitized_expected_contents = expected_contents.replace("\r", "")
+
+    self.assertEqual(
+        sanitized_actual_contents, sanitized_expected_contents, """
+pywrap_gradient_exclusions.cc needs to be updated.
+Please regenerate using:
+bazel run tensorflow/python/eager:gen_gradient_input_output_exclusions -- $PWD/tensorflow/python/eager/pywrap_gradient_exclusions.cc"""
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/graph_only_ops.py

@@ -0,0 +1,46 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Graph-only versions of a few op functions, for internal use only."""
+
+# Must be separate from array_ops to avoid a cyclic dependency.
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.framework import op_callbacks
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+
+
+def graph_placeholder(dtype, shape, name=None):
+  """Graph-only version of tf.compat.v1.placeholder(), for internal use only."""
+  dtype = dtype.base_dtype
+  dtype_value = attr_value_pb2.AttrValue(type=dtype.as_datatype_enum)
+  if isinstance(shape, (list, tuple)):
+    shape = tensor_shape.TensorShape(shape)
+  shape = attr_value_pb2.AttrValue(shape=shape.as_proto())
+  g = ops.get_default_graph()
+  attrs = {"dtype": dtype_value, "shape": shape}
+  op = g._create_op_internal(  # pylint: disable=protected-access
+      "Placeholder", [], [dtype], input_types=[],
+      attrs=attrs, name=name)
+  result, = op.outputs
+  if op_callbacks.should_invoke_op_callbacks():
+    # TODO(b/147670703): Once the special-op creation code paths
+    # are unified. Remove this `if` block.
+    callback_outputs = op_callbacks.invoke_op_callbacks(
+        "Placeholder", tuple(), attrs, tuple(op.outputs),
+        op_name=name, graph=g)
+    if callback_outputs is not None:
+      result, = callback_outputs
+  return result

tensorflow/python/eager/graph_only_ops_test.py

@@ -0,0 +1,40 @@
+# Copyright 2017 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 graph_only_ops."""
+
+import numpy as np
+
+from tensorflow.python.eager import graph_only_ops
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class GraphOnlyOpsTest(test_util.TensorFlowTestCase):
+
+  def testGraphPlaceholder(self):
+    with ops.Graph().as_default():
+      x_tf = graph_only_ops.graph_placeholder(dtypes.int32, shape=(1,))
+      y_tf = math_ops.square(x_tf)
+      with self.cached_session() as sess:
+        x = np.array([42])
+        y = sess.run(y_tf, feed_dict={x_tf: np.array([42])})
+        self.assertAllClose(np.square(x), y)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/imperative_grad.py

@@ -0,0 +1,73 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Code for backpropagation using the tape utilities."""
+
+import collections
+
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.ops.unconnected_gradients import UnconnectedGradients
+from tensorflow.python.util import compat
+
+VSpace = collections.namedtuple("VSpace", [
+    "aggregate_fn", "num_elements_fn", "zeros_fn", "ones_fn",
+    "zeros_like_fn", "ones_like_fn", "graph_shape_fn"
+])
+
+
+def imperative_grad(tape,
+                    target,
+                    sources,
+                    output_gradients=None,
+                    sources_raw=None,
+                    unconnected_gradients=UnconnectedGradients.NONE):
+  """Computes gradients from the imperatively defined tape on top of the stack.
+
+  Works by filtering the tape, computing how many downstream usages are of each
+  tensor and entry, and repeatedly applying backward functions until we have
+  gradients for all sources.
+
+  Args:
+   tape: the gradient tape which stores the trace.
+   target: either a Tensor or list of Tensors to be differentiated.
+   sources: list of Tensors for which we want gradients
+   output_gradients: if not None, a list of gradient provided for each Target,
+     or None if we are to use the target's computed downstream gradient.
+   sources_raw: if not None, a list of the source python objects from which the
+     sources were generated. Should have the same length as sources. Only needs
+     to be populated if unconnected_gradients is 'zero'.
+   unconnected_gradients: determines the value returned if the target and
+     sources are unconnected. When 'none' the value returned is None whereas
+     when 'zero' a zero tensor in the same shape as the sources is returned.
+
+  Returns:
+   the gradient wrt each of the sources.
+
+  Raises:
+    ValueError: if the arguments are invalid.
+    RuntimeError: if something goes wrong.
+  """
+  try:
+    unconnected_gradients = UnconnectedGradients(unconnected_gradients)
+  except ValueError:
+    raise ValueError(
+        "Unknown value for unconnected_gradients: %r" % unconnected_gradients)
+
+  return pywrap_tfe.TFE_Py_TapeGradient(
+      tape._tape,  # pylint: disable=protected-access
+      target,
+      sources,
+      output_gradients,
+      sources_raw,
+      compat.as_str(unconnected_gradients.value))

tensorflow/python/eager/lift_to_graph.py

@@ -0,0 +1,365 @@
+# Copyright 2018 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.
+# ==============================================================================
+# pylint: disable=unidiomatic-typecheck
+"""Utility to lift subgraphs."""
+
+import collections
+
+from tensorflow.python.framework import func_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import op_selector
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.util import compat
+from tensorflow.python.util import object_identity
+from tensorflow.python.util.tf_export import tf_export
+
+
+UnliftableError = op_selector.UnliftableError
+
+
+def _as_operation(op_or_tensor):
+  if isinstance(op_or_tensor, tensor_lib.Tensor):
+    return op_or_tensor.op
+  return op_or_tensor
+
+
+def _constant_inputs(op_or_tensor):
+  return all(_as_operation(i).type == u"Const"
+             and not _as_operation(i).control_inputs
+             for i in op_selector.graph_inputs(_as_operation(op_or_tensor)))
+
+
+# Represents an input to `copied_op` which must be updated once
+# `old_graph_tensor` has been copied.
+_InputMutation = collections.namedtuple(
+    "_InputMutation",
+    ["copied_op", "input_index", "old_graph_tensor"])
+
+
+# Represents a control input to `copied_op` which must be added once
+# `old_graph_op` has been copied.
+_ControlMutation = collections.namedtuple(
+    "_ControlMutation",
+    ["copied_op", "old_graph_op"])
+
+
+def _copy_non_source(op, graph, op_map, base_graph):
+  """Copy an op directly to a given graph.
+
+  Generally `op`'s inputs should already have been copied. If this is not the
+  case, for example with v1 while_loops, then `_copy_non_source` inserts
+  placeholders for the unavailable Tensors and returns a list of required
+  mutations.
+
+  Args:
+    op: The op to be copied.
+    graph: The destination graph.
+    op_map: A dict mapping ops and tensors in the old graph to the new one.
+    base_graph: The graph we're copying from, for any necessary functions.
+  Returns:
+    A tuple of (required_inputs, required_control_inputs):
+      required_inputs:
+        A list of `_InputMutation` tuples containing inputs to `copied_op` which
+        must be updated once `old_graph_tensor` has been copied.
+      required_control_inputs:
+        A list of `_ControlMutation` tuples containing control inputs to
+        `copied_op` which must be added once `old_graph_op` has been copied.
+  """
+  input_mutations = []
+  control_mutations = []
+  copied_inputs = []
+  for input_index, original_input in enumerate(op.inputs):
+    copied_input = op_map.get(original_input, None)
+    if copied_input is None:
+      # An input for this op is missing due to a loop in the graph. We'll insert
+      # a placeholder for now and return information about the required post-hoc
+      # mutation.
+      copied_input = array_ops.placeholder(
+          name="unused_control_flow_input",
+          shape=original_input.shape,
+          dtype=original_input.dtype)
+      input_mutations.append(
+          # `copied_op` is filled in below, after we've created it.
+          _InputMutation(copied_op=None,
+                         input_index=input_index,
+                         old_graph_tensor=original_input))
+    copied_inputs.append(copied_input)
+
+  copied_control_inputs = []
+  for original_control_input in op.control_inputs:
+    copied_control_input = op_map.get(original_control_input, None)
+    if copied_control_input is None:
+      control_mutations.append(
+          _ControlMutation(copied_op=None,
+                           old_graph_op=original_control_input))
+    else:
+      copied_control_inputs.append(copied_control_input)
+
+  # Don't copy over nodes with _tpu_replicate attribute. This attributed is used
+  # to signal that the op was built inside a tpu_replicate context; if we're
+  # lifting it to another graph we're similarly lifting it into another context.
+  with ops.control_dependencies(copied_control_inputs), ops.device(op.device):
+    # pylint: disable=protected-access
+    f = base_graph._functions.get(op.type, None)
+    if f is not None and compat.as_str(f.name) not in graph._functions:
+      f.add_to_graph(graph)
+    # pylint: enable=protected-access
+
+    # Create a new op in the destination graph if it doesn't exist before.
+    copied_op = graph.create_op(
+        op_type=op.type,
+        inputs=copied_inputs,
+        dtypes=[x.dtype for x in op.outputs],
+        attrs={
+            key: value for key, value in op.node_def.attr.items()
+            if not key.startswith("_class") and
+            not key.startswith("_tpu_replicate")
+        },  # b/128981532.
+        name=op.name)
+  op_map[op] = copied_op
+  for i, o in enumerate(op.outputs):
+    op_map[o] = copied_op.outputs[i]
+
+  return ([mutation._replace(copied_op=copied_op)
+           for mutation in input_mutations],
+          [mutation._replace(copied_op=copied_op)
+           for mutation in control_mutations])
+
+
+def _copy_source(s, graph, op_map, handle_captures, inverse_captures,
+                 base_graph):
+  """Create a source in a graph based on a Tensor from a different graph.
+
+  This function creates a placeholder analog of `s` in a graph with the
+  following behavior:
+
+  1) If s is a captured Tensor or Variable and handle_captures is set to True,
+     simply capture it in the new graph as well.
+
+  2) If s is a PlaceholderWithDefault whose default is a constant, preserve
+     said default in the new graph.
+
+  3) When applicable, copy resource variable metadata from `s` to the newly
+     created placeholder.
+
+  Args:
+    s: The source of interest.
+    graph: The destination graph.
+    op_map: A dict mapping ops and tensors in the old graph to the new one.
+    handle_captures: A boolean indicating whether to re-capture s in the new
+      graph or simply create a vanilla placeholder.
+    inverse_captures: A dict mapping s back to the Tensor or Variable that it
+      captures.
+    base_graph: The graph being copied from.
+  """
+  if handle_captures and s in inverse_captures:
+    copied_placeholder = graph.capture(inverse_captures[s], name=s.op.name)
+  elif s.op.type == "PlaceholderWithDefault" and _constant_inputs(s):
+    # Copy the default value to the graph.
+    default_value = s.op.inputs[0]
+    unavailable_inputs, unavailable_control_inputs = _copy_non_source(
+        op=default_value.op, graph=graph, op_map=op_map,
+        base_graph=base_graph)
+    if unavailable_inputs or unavailable_control_inputs:
+      raise AssertionError(
+          "Could not copy source node {} because it has inputs."
+          .format(default_value))
+
+    with ops.device(s.op.device):
+      copied_placeholder = array_ops.placeholder_with_default(
+          input=op_map[default_value], shape=s.shape, name=s.op.name)
+  else:
+    with ops.device(s.op.device):
+      copied_placeholder = array_ops.placeholder(
+          dtype=s.dtype, shape=s.shape, name=s.op.name)
+
+  base_handle = resource_variable_ops.get_resource_handle_data(s)
+  if base_handle.shape_and_type:
+    resource_variable_ops._set_handle_shapes_and_types(  # pylint: disable=protected-access
+        copied_placeholder,
+        base_handle,
+        graph_mode=True)
+
+  op_map[s] = copied_placeholder
+  # Add an entry for the op of the source tensor so that if there are any nodes
+  # depending on that op via control dependencies it can work correctly.
+  op_map[s.op] = copied_placeholder.op
+
+
+@tf_export("__internal__.lift_to_graph", v1=[])
+def lift_to_graph(tensors,
+                  graph,
+                  sources=None,
+                  disallowed_placeholders=None,
+                  add_sources=False,
+                  handle_captures=False,
+                  base_graph=None,
+                  op_map=None):
+  """Copies the tensor and all its inputs recursively to the outer graph.
+
+  Args:
+    tensors: The Tensors to lift.
+    graph: The graph to lift to.
+    sources: Optional sequence of nodes to start from. If omitted the whole
+      subgraph which feeds into `init_tensor` is lifted.
+    disallowed_placeholders: An optional set of ops which may not appear in the
+      lifted graph. Defaults to all placeholders.
+    add_sources: A boolean indicating whether placeholders which are not in
+      sources should be allowed.
+    handle_captures: A boolean indicating whether to re-capture s in the new
+      graph or simply create a vanilla placeholder.
+    base_graph: The graph from which to lift ops. This will be inferred if not
+      specified.
+    op_map: A map contains all the existing nodes that have been lifted to the
+      destination graph, so they won't be lifted and copied again.
+
+  Returns:
+    A mapping from ops in the current default graph to ops in `graph`.
+
+  Raises:
+    UnliftableError: If a placeholder blocks lifting.
+  """
+  variable_init_tensors = []
+  init_tensors = []
+  for tensor in tensors:
+    if isinstance(tensor, resource_variable_ops.ResourceVariable):
+      variable_init_tensors.append(tensor)
+    else:
+      init_tensors.append(tensor)
+  base_graph = base_graph or init_tensors[0].graph
+  op_map = op_map or object_identity.ObjectIdentityDictionary()
+
+  # Check that the initializer does not depend on any placeholders.
+  sources = object_identity.ObjectIdentitySet(sources or [])
+  visited_ops = set(x.op for x in sources)
+  op_outputs = collections.defaultdict(set)
+
+  # First we extract the subgraph between init_tensors and sources.
+  for init_tensor in init_tensors:
+    sources.update(op_selector.map_subgraph(
+        init_tensor=init_tensor,
+        sources=sources,
+        disallowed_placeholders=disallowed_placeholders,
+        visited_ops=visited_ops,
+        op_outputs=op_outputs,
+        add_sources=add_sources))
+
+  # Try to topologically sort the nodes we've extracted. Now we know how many of
+  # their outputs are part of this subgraph.
+  ops_to_copy = []
+  marked_ops = set([])
+  ops_to_visit = [_as_operation(t) for t in init_tensors
+                  if not op_outputs[_as_operation(t)]]
+  unvisited_ops = set(ops_to_visit)
+  while unvisited_ops:
+    while ops_to_visit:
+      op = ops_to_visit.pop()
+      if op in marked_ops:
+        continue
+      marked_ops.add(op)
+      ops_to_copy.append(op)
+      for inp in op_selector.graph_inputs(op):
+        # Don't lift the TPUReplicateMetadata nodes out of the function, because
+        # it has no registered kernels.
+        if inp.type == "TPUReplicateMetadata":
+          continue
+        unvisited_ops.add(inp)
+        if (all(x in marked_ops for x in op_outputs[inp]) and
+            inp not in sources):
+          ops_to_visit.append(inp)
+    unvisited_ops.difference_update(marked_ops)
+    if unvisited_ops:
+      # `unvisited_ops` should only have elements if the graph has a loop. In
+      # this case we want to keep copying and there's no topological ordering;
+      # we'll do ugly post-hoc mutations instead.
+      ops_to_visit.append(next(iter(unvisited_ops)))
+
+  # When the topological sort fails due to loops, it can result in exceptions
+  # later when copying a node which inputs haven't been copied yet. We can
+  # improve that pseudo-topological order slightly by putting the ops without
+  # inputs, such as constants, at the start of the topological order (i.e at
+  # the end of ops_to_copy).
+  ops_to_copy.sort(key=(lambda op: len(op_selector.graph_inputs(op)) == 0))
+
+  # When lifting from one FuncGraph to another, we will need to capture the
+  # relevant tensors as well.
+  captures = []
+  inverse_captures = object_identity.ObjectIdentityDictionary()
+  internal_captures = []
+  if (isinstance(base_graph, func_graph.FuncGraph) and
+      isinstance(graph, func_graph.FuncGraph)):
+    captures = base_graph.captures
+    for external_capture, internal_capture in captures:
+      inverse_captures[internal_capture] = external_capture
+    internal_captures = base_graph.internal_captures
+
+  # ops_to_copy now holds a reverse topologically sorted list of ops which
+  # ends in the initializer. We copy those to the outermost graph and
+  # build the initialization op there.
+  with graph.as_default():
+    for i in variable_init_tensors:
+      op_map[i] = i
+    source_ops = set()
+    # Add the sources in the same order as the original graph.
+    for s in internal_captures:
+      if s in sources:
+        sources.remove(s)
+        source_ops.add(s.op)
+        _copy_source(
+            s=s,
+            graph=graph,
+            op_map=op_map,
+            handle_captures=handle_captures,
+            inverse_captures=inverse_captures,
+            base_graph=base_graph)
+    for s in sources:
+      source_ops.add(s.op)
+      _copy_source(
+          s=s,
+          graph=graph,
+          op_map=op_map,
+          handle_captures=handle_captures,
+          inverse_captures=inverse_captures,
+          base_graph=base_graph)
+
+    input_mutations = []
+    control_mutations = []
+    for op in reversed(ops_to_copy):
+      if op in source_ops or op in op_map:
+        continue
+      new_input_mutations, new_control_mutations = _copy_non_source(
+          op=op, graph=graph, op_map=op_map, base_graph=base_graph)
+      input_mutations.extend(new_input_mutations)
+      control_mutations.extend(new_control_mutations)
+
+    # Mutate the new graph to insert any loops which existed in the source
+    # graph due to v1 while_loops.
+    #
+    # pylint: disable=protected-access
+    with graph._mutation_lock():
+      for mutation in input_mutations:
+        mutation.copied_op._update_input(
+            mutation.input_index, op_map[mutation.old_graph_tensor])
+      for mutation in control_mutations:
+        # Don't lift the TPUReplicateMetadata nodes out of the function, because
+        # it has no registered kernels.
+        if mutation.old_graph_op.type == "TPUReplicateMetadata":
+          continue
+        mutation.copied_op._add_control_input(op_map[mutation.old_graph_op])
+    # pylint: enable=protected-access
+
+    return op_map

tensorflow/python/eager/lift_to_graph_test.py

@@ -0,0 +1,88 @@
+# Copyright 2019 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 lift_to_graph."""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import lift_to_graph
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import func_graph
+from tensorflow.python.framework import ops as framework_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.util import compat
+
+
+class LiftToGraphTest(test.TestCase):
+
+  def testCaptureOrdering(self):
+    v1 = resource_variable_ops.ResourceVariable(1.0)
+    v2 = resource_variable_ops.ResourceVariable(2.0)
+    v3 = resource_variable_ops.ResourceVariable(3.0)
+
+    @def_function.function
+    def fn():
+      return v1 + v2 + v3
+
+    concrete_fn = fn.get_concrete_function()
+    original_captures = concrete_fn.graph.internal_captures
+    outputs = concrete_fn.graph.outputs
+
+    for _ in range(100):
+      g = func_graph.FuncGraph('lifted')
+
+      lift_to_graph.lift_to_graph(
+          outputs, g, add_sources=True, handle_captures=True)
+      lifted_captures = g.internal_captures
+      self.assertLen(lifted_captures, 3)
+      for original, lifted in zip(original_captures, lifted_captures):
+        self.assertEqual(original.name, lifted.name)
+
+  def testClassAttrsRemoved(self):
+    """Tests that _class attrs (from colocate_with()) are removed."""
+    @def_function.function
+    def fn():
+      two = constant_op.constant(2.0, name='two')
+      ten = constant_op.constant(10.0, name='ten')
+      twenty = math_ops.multiply(two, ten, name='twenty')
+      three = constant_op.constant(3.0, name='three')
+      with framework_ops.colocate_with(twenty):
+        thirty = math_ops.multiply(three, ten, name='thirty')
+      return ten, twenty, thirty
+
+    concrete_fn = fn.get_concrete_function()
+    self.assertItemsEqual(  # Before lifting, 'fn' has colocation attrs.
+        concrete_fn.graph.get_operation_by_name('thirty').colocation_groups(),
+        [compat.as_bytes('loc:@twenty')])
+    thirty_out = concrete_fn.graph.outputs[2]
+
+    g = func_graph.FuncGraph('lifted')
+    lift_to_graph.lift_to_graph([thirty_out], g)
+
+    # After lifting, colocation attrs are gone.
+    ops = g.get_operations()
+    self.assertItemsEqual([op.name for op in ops],
+                          ['three', 'ten', 'thirty',  # Lifted from `fn` body.
+                           thirty_out.op.name])  # Wrapper for output.
+    for op in ops:
+      with self.assertRaises(ValueError):
+        class_attr = op.get_attr('_class')  # Expected not to exist.
+        print('Unexpected class_attr', class_attr, 'on', op.name)
+      self.assertItemsEqual(op.colocation_groups(),  # Expect default self-ref.
+                            [compat.as_bytes('loc:@%s' % op.name)])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/memory_tests/BUILD

@@ -0,0 +1,67 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+# NOTE: Do not add sharding to these tests. If tests run concurrently, they
+# seem to confuse the memory_profiler, and the tests begin to flake. Add new
+# test files as needed.
+
+py_library(
+    name = "memory_test_util",
+    srcs = ["memory_test_util.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = ["//tensorflow/python/eager:context"],
+)
+
+cuda_py_strict_test(
+    name = "memory_test",
+    size = "medium",
+    srcs = ["memory_test.py"],
+    tags = [
+        "manual",
+        "no_oss",
+        "notap",  #TODO(b/140640597): this test is flaky at the moment
+        "optonly",  # The test is too slow in non-opt mode
+    ],
+    # TODO(b/140065350): Re-enable
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":memory_test_util",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:gradients",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "remote_memory_test",
+    size = "medium",
+    srcs = ["remote_memory_test.py"],
+    tags = [
+        "no_gpu",  # TODO(b/168058741): Enable the test for GPU
+        "optonly",  # The test is too slow in non-opt mode
+    ],
+    xla_enable_strict_auto_jit = False,  # b/140261762
+    deps = [
+        ":memory_test_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:remote",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/training:server_lib",
+    ],
+)

tensorflow/python/eager/memory_tests/memory_test.py

@@ -0,0 +1,117 @@
+# Copyright 2018 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 memory leaks in eager execution.
+
+It is possible that this test suite will eventually become flaky due to taking
+too long to run (since the tests iterate many times), but for now they are
+helpful for finding memory leaks since not all PyObject leaks are found by
+introspection (test_util decorators). Please be careful adding new tests here.
+"""
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.eager.memory_tests import memory_test_util
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients as gradient_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.variables import Variable
+
+
+class MemoryTest(test.TestCase):
+
+  def testMemoryLeakAnonymousVariable(self):
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    def f():
+      inputs = Variable(array_ops.zeros([32, 100], dtypes.float32))
+      del inputs
+
+    memory_test_util.assert_no_leak(
+        f, num_iters=10000, increase_threshold_absolute_mb=10)
+
+  def testMemoryLeakInFunction(self):
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    def f():
+
+      @def_function.function
+      def graph(x):
+        return x * x + x
+
+      graph(constant_op.constant(42))
+
+    memory_test_util.assert_no_leak(
+        f, num_iters=1000, increase_threshold_absolute_mb=30)
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testNestedFunctionsDeleted(self):
+
+    @def_function.function
+    def f(x):
+
+      @def_function.function
+      def my_sin(x):
+        return math_ops.sin(x)
+
+      return my_sin(x)
+
+    x = constant_op.constant(1.)
+    with backprop.GradientTape() as t1:
+      t1.watch(x)
+      with backprop.GradientTape() as t2:
+        t2.watch(x)
+        y = f(x)
+      dy_dx = t2.gradient(y, x)
+    dy2_dx2 = t1.gradient(dy_dx, x)
+
+    self.assertAllClose(0.84147096, y.numpy())  # sin(1.)
+    self.assertAllClose(0.54030230, dy_dx.numpy())  # cos(1.)
+    self.assertAllClose(-0.84147096, dy2_dx2.numpy())  # -sin(1.)
+
+  def testMemoryLeakInGlobalGradientRegistry(self):
+    # Past leak: b/139819011
+
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    def f():
+
+      @def_function.function(autograph=False)
+      def graph(x):
+
+        @def_function.function(autograph=False)
+        def cubed(a):
+          return a * a * a
+
+        y = cubed(x)
+        # To ensure deleting the function does not affect the gradient
+        # computation.
+        del cubed
+        return gradient_ops.gradients(gradient_ops.gradients(y, x), x)
+
+      return graph(constant_op.constant(1.5))[0].numpy()
+
+    memory_test_util.assert_no_leak(
+        f, num_iters=300, increase_threshold_absolute_mb=50)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/memory_tests/memory_test_util.py

@@ -0,0 +1,73 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Utils for memory tests."""
+
+import collections
+import gc
+import time
+
+from tensorflow.python.eager import context
+
+# memory_profiler might not be available in the OSS version of TensorFlow.
+try:
+  import memory_profiler  # pylint:disable=g-import-not-at-top
+except ImportError:
+  memory_profiler = None
+
+
+def _instance_count_by_class():
+  counter = collections.Counter()
+
+  for obj in gc.get_objects():
+    try:
+      counter[obj.__class__.__name__] += 1
+    except Exception:  # pylint:disable=broad-except
+      pass
+
+  return counter
+
+
+def assert_no_leak(f, num_iters=100000, increase_threshold_absolute_mb=25):
+  """Assert memory usage doesn't increase beyond given threshold for f."""
+
+  with context.eager_mode():
+    # Warm up.
+    f()
+
+    # Wait for background threads to start up and take over memory.
+    # FIXME: The nature of this test leaves few other options. Maybe there
+    # is a better way to do this.
+    time.sleep(4)
+
+    gc.collect()
+    initial = memory_profiler.memory_usage(-1)[0]
+    instance_count_by_class_before = _instance_count_by_class()
+
+    for _ in range(num_iters):
+      f()
+
+    gc.collect()
+    increase = memory_profiler.memory_usage(-1)[0] - initial
+
+    assert increase < increase_threshold_absolute_mb, (
+        "Increase is too high. Initial memory usage: %f MB. Increase: %f MB. "
+        "Maximum allowed increase: %f MB. "
+        "Instance count diff before/after: %s") % (
+            initial, increase, increase_threshold_absolute_mb,
+            _instance_count_by_class() - instance_count_by_class_before)
+
+
+def memory_profiler_is_available():
+  return memory_profiler is not None

tensorflow/python/eager/memory_tests/remote_memory_test.py

@@ -0,0 +1,60 @@
+# Copyright 2019 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 memory leaks in remote eager execution."""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.eager.memory_tests import memory_test_util
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.training import server_lib
+
+
+class RemoteWorkerMemoryTest(test.TestCase):
+
+  def __init__(self, method):
+    super(RemoteWorkerMemoryTest, self).__init__(method)
+
+    # used for remote worker tests
+    self._cached_server = server_lib.Server.create_local_server()
+    self._cached_server_target = self._cached_server.target[len("grpc://"):]
+
+  def testMemoryLeakInLocalCopy(self):
+    if not memory_test_util.memory_profiler_is_available():
+      self.skipTest("memory_profiler required to run this test")
+
+    remote.connect_to_remote_host(self._cached_server_target)
+
+    # Run a function locally with the input on a remote worker and ensure we
+    # do not leak a reference to the remote tensor.
+
+    @def_function.function
+    def local_func(i):
+      return i
+
+    def func():
+      with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+        x = array_ops.zeros([1000, 1000], dtypes.int32)
+
+      local_func(x)
+
+    memory_test_util.assert_no_leak(
+        func, num_iters=100, increase_threshold_absolute_mb=50)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/monitoring.py

@@ -0,0 +1,542 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""TensorFlow monitoring APIs."""
+
+import collections
+import functools
+import time
+
+from tensorflow.core.framework import summary_pb2
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.client import pywrap_tf_session
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.util import compat
+from tensorflow.python.util.tf_export import tf_export
+
+_MetricMethod = collections.namedtuple('MetricMethod', 'create delete get_cell')
+_counter_methods = [
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewCounter0,
+        delete=pywrap_tfe.TFE_MonitoringDeleteCounter0,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellCounter0),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewCounter1,
+        delete=pywrap_tfe.TFE_MonitoringDeleteCounter1,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellCounter1),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewCounter2,
+        delete=pywrap_tfe.TFE_MonitoringDeleteCounter2,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellCounter2),
+]
+_int_gauge_methods = [
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewIntGauge0,
+        delete=pywrap_tfe.TFE_MonitoringDeleteIntGauge0,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellIntGauge0),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewIntGauge1,
+        delete=pywrap_tfe.TFE_MonitoringDeleteIntGauge1,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellIntGauge1),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewIntGauge2,
+        delete=pywrap_tfe.TFE_MonitoringDeleteIntGauge2,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellIntGauge2),
+]
+_string_gauge_methods = [
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewStringGauge0,
+        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge0,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge0),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewStringGauge1,
+        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge1,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge1),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewStringGauge2,
+        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge2,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge2),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewStringGauge3,
+        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge3,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge3),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewStringGauge4,
+        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge4,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge4),
+]
+_bool_gauge_methods = [
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewBoolGauge0,
+        delete=pywrap_tfe.TFE_MonitoringDeleteBoolGauge0,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellBoolGauge0),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewBoolGauge1,
+        delete=pywrap_tfe.TFE_MonitoringDeleteBoolGauge1,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellBoolGauge1),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewBoolGauge2,
+        delete=pywrap_tfe.TFE_MonitoringDeleteBoolGauge2,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellBoolGauge2),
+]
+_sampler_methods = [
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewSampler0,
+        delete=pywrap_tfe.TFE_MonitoringDeleteSampler0,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellSampler0),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewSampler1,
+        delete=pywrap_tfe.TFE_MonitoringDeleteSampler1,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellSampler1),
+    _MetricMethod(
+        create=pywrap_tfe.TFE_MonitoringNewSampler2,
+        delete=pywrap_tfe.TFE_MonitoringDeleteSampler2,
+        get_cell=pywrap_tfe.TFE_MonitoringGetCellSampler2),
+]
+
+
+class Metric(object):
+  """The base class of metric."""
+
+  __slots__ = ["_metric", "_metric_name", "_metric_methods", "_label_length"]
+
+  def __init__(self, metric_name, metric_methods, label_length, *args):
+    """Creates a new metric.
+
+    Args:
+      metric_name: name of the metric class.
+      metric_methods: list of swig metric methods.
+      label_length: length of label args.
+      *args: the arguments to call create method.
+    """
+    self._metric_name = metric_name
+    self._metric_methods = metric_methods
+    self._label_length = label_length
+
+    if label_length >= len(self._metric_methods):
+      raise ValueError('Cannot create {} metric with label >= {}'.format(
+          self._metric_name, len(self._metric_methods)))
+
+    self._metric = self._metric_methods[self._label_length].create(*args)
+
+  def __del__(self):
+    try:
+      deleter = self._metric_methods[self._label_length].delete
+      metric = self._metric
+    except AttributeError:
+      return
+
+    if deleter is not None:
+      deleter(metric)
+
+  def get_cell(self, *labels):
+    """Retrieves the cell."""
+    if len(labels) != self._label_length:
+      raise ValueError('The {} expects taking {} labels'.format(
+          self._metric_name, self._label_length))
+    return self._metric_methods[self._label_length].get_cell(
+        self._metric, *labels)
+
+
+class CounterCell(object):
+  """CounterCell stores each value of a Counter."""
+
+  __slots__ = ["_cell"]
+
+  def __init__(self, cell):
+    """Creates a new CounterCell.
+
+    Args:
+      cell: A c pointer of TFE_MonitoringCounterCell.
+    """
+    self._cell = cell
+
+  def increase_by(self, value):
+    """Atomically increments the value.
+
+    Args:
+      value: non-negative value.
+    """
+    pywrap_tfe.TFE_MonitoringCounterCellIncrementBy(self._cell, value)
+
+  def value(self):
+    """Retrieves the current value."""
+    return pywrap_tfe.TFE_MonitoringCounterCellValue(self._cell)
+
+
+class Counter(Metric):
+  """A stateful class for updating a cumulative integer metric.
+
+  This class encapsulates a set of values (or a single value for a label-less
+  metric). Each value is identified by a tuple of labels. The class allows the
+  user to increment each value.
+  """
+
+  __slots__ = []
+
+  def __init__(self, name, description, *labels):
+    """Creates a new Counter.
+
+    Args:
+      name: name of the new metric.
+      description: description of the new metric.
+      *labels: The label list of the new metric.
+    """
+    super(Counter, self).__init__('Counter', _counter_methods, len(labels),
+                                  name, description, *labels)
+
+  def get_cell(self, *labels):
+    """Retrieves the cell."""
+    return CounterCell(super(Counter, self).get_cell(*labels))
+
+
+class IntGaugeCell(object):
+  """A single integer value stored in an `IntGauge`."""
+
+  __slots__ = ["_cell"]
+
+  def __init__(self, cell):
+    """Creates a new IntGaugeCell.
+
+    Args:
+      cell: A c pointer of TFE_MonitoringIntGaugeCell.
+    """
+    self._cell = cell
+
+  def set(self, value):
+    """Atomically set the value.
+
+    Args:
+      value: integer value.
+    """
+    pywrap_tfe.TFE_MonitoringIntGaugeCellSet(self._cell, value)
+
+  def value(self):
+    """Retrieves the current value."""
+    return pywrap_tfe.TFE_MonitoringIntGaugeCellValue(self._cell)
+
+
+class IntGauge(Metric):
+  """A stateful class for updating a gauge-like integer metric.
+
+  This class encapsulates a set of integer values (or a single value for a
+  label-less metric). Each value is identified by a tuple of labels. The class
+  allows the user to set each value.
+  """
+
+  __slots__ = []
+
+  def __init__(self, name, description, *labels):
+    """Creates a new IntGauge.
+
+    Args:
+      name: name of the new metric.
+      description: description of the new metric.
+      *labels: The label list of the new metric.
+    """
+    super(IntGauge, self).__init__('IntGauge', _int_gauge_methods, len(labels),
+                                   name, description, *labels)
+
+  def get_cell(self, *labels):
+    """Retrieves the cell."""
+    return IntGaugeCell(super(IntGauge, self).get_cell(*labels))
+
+
+class StringGaugeCell(object):
+  """A single string value stored in an `StringGauge`."""
+
+  __slots__ = ["_cell"]
+
+  def __init__(self, cell):
+    """Creates a new StringGaugeCell.
+
+    Args:
+      cell: A c pointer of TFE_MonitoringStringGaugeCell.
+    """
+    self._cell = cell
+
+  def set(self, value):
+    """Atomically set the value.
+
+    Args:
+      value: string value.
+    """
+    pywrap_tfe.TFE_MonitoringStringGaugeCellSet(self._cell, value)
+
+  def value(self):
+    """Retrieves the current value."""
+    with c_api_util.tf_buffer() as buffer_:
+      pywrap_tfe.TFE_MonitoringStringGaugeCellValue(self._cell, buffer_)
+      value = pywrap_tf_session.TF_GetBuffer(buffer_).decode('utf-8')
+    return value
+
+
+class StringGauge(Metric):
+  """A stateful class for updating a gauge-like string metric.
+
+  This class encapsulates a set of string values (or a single value for a
+  label-less metric). Each value is identified by a tuple of labels. The class
+  allows the user to set each value.
+  """
+
+  __slots__ = []
+
+  def __init__(self, name, description, *labels):
+    """Creates a new StringGauge.
+
+    Args:
+      name: name of the new metric.
+      description: description of the new metric.
+      *labels: The label list of the new metric.
+    """
+    super(StringGauge, self).__init__('StringGauge', _string_gauge_methods,
+                                      len(labels), name, description, *labels)
+
+  def get_cell(self, *labels):
+    """Retrieves the cell."""
+    return StringGaugeCell(super(StringGauge, self).get_cell(*labels))
+
+
+class BoolGaugeCell(object):
+  """A single boolean value stored in an `BoolGauge`."""
+
+  __slots__ = ["_cell"]
+
+  def __init__(self, cell):
+    """Creates a new BoolGaugeCell.
+
+    Args:
+      cell: A c pointer of TFE_MonitoringBoolGaugeCell.
+    """
+    self._cell = cell
+
+  def set(self, value):
+    """Atomically set the value.
+
+    Args:
+      value: bool value.
+    """
+    pywrap_tfe.TFE_MonitoringBoolGaugeCellSet(self._cell, value)
+
+  def value(self):
+    """Retrieves the current value."""
+    return pywrap_tfe.TFE_MonitoringBoolGaugeCellValue(self._cell)
+
+
+@tf_export("__internal__.monitoring.BoolGauge", v1=[])
+class BoolGauge(Metric):
+  """A stateful class for updating a gauge-like bool metric.
+
+  This class encapsulates a set of boolean values (or a single value for a
+  label-less metric). Each value is identified by a tuple of labels. The class
+  allows the user to set each value.
+  """
+
+  __slots__ = []
+
+  def __init__(self, name, description, *labels):
+    """Creates a new BoolGauge.
+
+    Args:
+      name: name of the new metric.
+      description: description of the new metric.
+      *labels: The label list of the new metric.
+    """
+    super(BoolGauge, self).__init__('BoolGauge', _bool_gauge_methods,
+                                    len(labels), name, description, *labels)
+
+  def get_cell(self, *labels):
+    """Retrieves the cell."""
+    return BoolGaugeCell(super(BoolGauge, self).get_cell(*labels))
+
+
+class SamplerCell(object):
+  """SamplerCell stores each value of a Sampler."""
+
+  __slots__ = ["_cell"]
+
+  def __init__(self, cell):
+    """Creates a new SamplerCell.
+
+    Args:
+      cell: A c pointer of TFE_MonitoringSamplerCell.
+    """
+    self._cell = cell
+
+  def add(self, value):
+    """Atomically add a sample.
+
+    Args:
+      value: float value.
+    """
+    pywrap_tfe.TFE_MonitoringSamplerCellAdd(self._cell, value)
+
+  def value(self):
+    """Retrieves the current distribution of samples.
+
+    Returns:
+      A HistogramProto describing the distribution of samples.
+    """
+    with c_api_util.tf_buffer() as buffer_:
+      pywrap_tfe.TFE_MonitoringSamplerCellValue(self._cell, buffer_)
+      proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
+    histogram_proto = summary_pb2.HistogramProto()
+    histogram_proto.ParseFromString(compat.as_bytes(proto_data))
+    return histogram_proto
+
+
+class Buckets(object):
+  """Bucketing strategies for the samplers."""
+
+  __slots__ = ["buckets"]
+
+  def __init__(self, buckets):
+    """Creates a new Buckets.
+
+    Args:
+      buckets: A c pointer of TFE_MonitoringBuckets.
+    """
+    self.buckets = buckets
+
+  def __del__(self):
+    pywrap_tfe.TFE_MonitoringDeleteBuckets(self.buckets)
+
+
+class ExponentialBuckets(Buckets):
+  """Exponential bucketing strategy.
+
+  Sets up buckets of the form:
+      [-DBL_MAX, ..., scale * growth^i,
+       scale * growth_factor^(i + 1), ..., DBL_MAX].
+  """
+
+  __slots__ = []
+
+  def __init__(self, scale, growth_factor, bucket_count):
+    """Creates a new exponential Buckets.
+
+    Args:
+      scale: float
+      growth_factor: float
+      bucket_count: integer
+    """
+    super(ExponentialBuckets, self).__init__(
+        pywrap_tfe.TFE_MonitoringNewExponentialBuckets(scale, growth_factor,
+                                                       bucket_count))
+
+
+class Sampler(Metric):
+  """A stateful class for updating a cumulative histogram metric.
+
+  This class encapsulates a set of histograms (or a single histogram for a
+  label-less metric) configured with a list of increasing bucket boundaries.
+  Each histogram is identified by a tuple of labels. The class allows the
+  user to add a sample to each histogram value.
+  """
+
+  __slots__ = []
+
+  def __init__(self, name, buckets, description, *labels):
+    """Creates a new Sampler.
+
+    Args:
+      name: name of the new metric.
+      buckets: bucketing strategy of the new metric.
+      description: description of the new metric.
+      *labels: The label list of the new metric.
+    """
+    super(Sampler, self).__init__('Sampler', _sampler_methods, len(labels),
+                                  name, buckets.buckets, description, *labels)
+
+  def get_cell(self, *labels):
+    """Retrieves the cell."""
+    return SamplerCell(super(Sampler, self).get_cell(*labels))
+
+
+# Keeping track of current MonitoredTimer sections to prevent repetitive
+# counting.
+MonitoredTimerSections = []
+
+
+class MonitoredTimer(object):
+  """A context manager to measure the walltime and increment a Counter cell."""
+
+  __slots__ = [
+      "cell",
+      "t",
+      "monitored_section_name",
+      "_counting",
+      "_avoid_repetitive_counting",
+  ]
+
+  def __init__(
+      self, cell, monitored_section_name=None, avoid_repetitive_counting=False
+  ):
+    """Creates a new MonitoredTimer.
+
+    Args:
+      cell: the cell associated with the time metric that will be inremented.
+      monitored_section_name: name of action being monitored here.
+      avoid_repetitive_counting: when set to True, if already in a monitored
+        timer section with the same monitored_section_name, skip counting.
+    """
+    self.cell = cell
+    self.monitored_section_name = monitored_section_name
+    self._avoid_repetitive_counting = avoid_repetitive_counting
+    self._counting = True
+
+  def __enter__(self):
+    if (
+        self._avoid_repetitive_counting
+        and self.monitored_section_name
+        and self.monitored_section_name in MonitoredTimerSections
+    ):
+      self._counting = False
+      return self
+
+    self.t = time.time()
+    if self.monitored_section_name:
+      MonitoredTimerSections.append(self.monitored_section_name)
+
+    return self
+
+  def __exit__(self, exception_type, exception_value, traceback):
+    del exception_type, exception_value, traceback
+    if self._counting:
+      micro_seconds = (time.time() - self.t) * 1000000
+      self.cell.increase_by(int(micro_seconds))
+      if self.monitored_section_name:
+        MonitoredTimerSections.remove(self.monitored_section_name)
+
+
+def monitored_timer(cell):
+  """A function decorator for adding MonitoredTimer support.
+
+  Args:
+    cell: the cell associated with the time metric that will be inremented.
+  Returns:
+    A decorator that measure the function runtime and increment the specified
+    counter cell.
+  """
+
+  def actual_decorator(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+      with MonitoredTimer(cell):
+        return func(*args, **kwargs)
+
+    return wrapper
+
+  return actual_decorator

tensorflow/python/eager/monitoring_test.py

@@ -0,0 +1,159 @@
+# Copyright 2018 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 monitoring."""
+
+import time
+
+from tensorflow.python.eager import monitoring
+from tensorflow.python.eager import test
+from tensorflow.python.framework import test_util
+
+
+class MonitoringTest(test_util.TensorFlowTestCase):
+
+  def test_counter(self):
+    counter = monitoring.Counter('test/counter', 'test counter')
+    counter.get_cell().increase_by(1)
+    self.assertEqual(counter.get_cell().value(), 1)
+    counter.get_cell().increase_by(5)
+    self.assertEqual(counter.get_cell().value(), 6)
+
+  def test_multiple_counters(self):
+    counter1 = monitoring.Counter('test/counter1', 'test counter', 'label1')
+    counter1.get_cell('foo').increase_by(1)
+    self.assertEqual(counter1.get_cell('foo').value(), 1)
+    counter2 = monitoring.Counter('test/counter2', 'test counter', 'label1',
+                                  'label2')
+    counter2.get_cell('foo', 'bar').increase_by(5)
+    self.assertEqual(counter2.get_cell('foo', 'bar').value(), 5)
+
+  def test_same_counter(self):
+    counter1 = monitoring.Counter('test/same_counter', 'test counter')  # pylint: disable=unused-variable
+    counter2 = monitoring.Counter('test/same_counter', 'test counter')  # pylint: disable=unused-variable
+
+  def test_int_gauge(self):
+    gauge = monitoring.IntGauge('test/gauge', 'test gauge')
+    gauge.get_cell().set(1)
+    self.assertEqual(gauge.get_cell().value(), 1)
+    gauge.get_cell().set(5)
+    self.assertEqual(gauge.get_cell().value(), 5)
+
+    gauge1 = monitoring.IntGauge('test/gauge1', 'test gauge1', 'label1')
+    gauge1.get_cell('foo').set(2)
+    self.assertEqual(gauge1.get_cell('foo').value(), 2)
+
+  def test_string_gauge(self):
+    gauge = monitoring.StringGauge('test/gauge', 'test gauge')
+    gauge.get_cell().set('left')
+    self.assertEqual(gauge.get_cell().value(), 'left')
+    gauge.get_cell().set('right')
+    self.assertEqual(gauge.get_cell().value(), 'right')
+
+    gauge1 = monitoring.StringGauge('test/gauge1', 'test gauge1', 'label1')
+    gauge1.get_cell('foo').set('start')
+    self.assertEqual(gauge1.get_cell('foo').value(), 'start')
+
+  def test_bool_gauge(self):
+    gauge = monitoring.BoolGauge('test/gauge', 'test gauge')
+    gauge.get_cell().set(True)
+    self.assertTrue(gauge.get_cell().value())
+    gauge.get_cell().set(False)
+    self.assertFalse(gauge.get_cell().value())
+
+    gauge1 = monitoring.BoolGauge('test/gauge1', 'test gauge1', 'label1')
+    gauge1.get_cell('foo').set(True)
+    self.assertTrue(gauge1.get_cell('foo').value())
+
+  def test_sampler(self):
+    buckets = monitoring.ExponentialBuckets(1.0, 2.0, 2)
+    sampler = monitoring.Sampler('test/sampler', buckets, 'test sampler')
+    sampler.get_cell().add(1.0)
+    sampler.get_cell().add(5.0)
+    histogram_proto = sampler.get_cell().value()
+    self.assertEqual(histogram_proto.min, 1.0)
+    self.assertEqual(histogram_proto.num, 2.0)
+    self.assertEqual(histogram_proto.sum, 6.0)
+
+    sampler1 = monitoring.Sampler('test/sampler1', buckets, 'test sampler',
+                                  'label1')
+    sampler1.get_cell('foo').add(2.0)
+    sampler1.get_cell('foo').add(4.0)
+    sampler1.get_cell('bar').add(8.0)
+    histogram_proto1 = sampler1.get_cell('foo').value()
+    self.assertEqual(histogram_proto1.max, 4.0)
+    self.assertEqual(histogram_proto1.num, 2.0)
+    self.assertEqual(histogram_proto1.sum, 6.0)
+
+  def test_context_manager(self):
+    counter = monitoring.Counter('test/ctxmgr', 'test context manager', 'slot')
+    with monitoring.MonitoredTimer(counter.get_cell('long')):
+      time.sleep(0.01)
+      with monitoring.MonitoredTimer(counter.get_cell('short')):
+        time.sleep(0.01)
+    self.assertGreater(
+        counter.get_cell('long').value(), counter.get_cell('short').value()
+    )
+
+  def test_monitored_timer_tracker(self):
+    counter = monitoring.Counter('test/ctxmgr', 'test context manager', 'slot')
+    counter2 = monitoring.Counter('test/ctxmgr2', 'slot')
+    with monitoring.MonitoredTimer(counter.get_cell('long'), 'counter'):
+      time.sleep(0.01)
+      self.assertIn('counter', monitoring.MonitoredTimerSections)
+      with monitoring.MonitoredTimer(counter2.get_cell(), 'counter2'):
+        time.sleep(0.01)
+        self.assertIn('counter', monitoring.MonitoredTimerSections)
+        self.assertIn('counter2', monitoring.MonitoredTimerSections)
+        with monitoring.MonitoredTimer(counter.get_cell('long'), 'counter'):
+          time.sleep(0.01)
+      self.assertNotIn('counter2', monitoring.MonitoredTimerSections)
+    self.assertGreater(
+        counter.get_cell('long').value(), counter.get_cell('short').value()
+    )
+    self.assertGreater(counter2.get_cell().value(), 0)
+
+  def test_repetitive_monitored_timer(self):
+    counter = monitoring.Counter('test/ctxmgr', 'test context manager')
+    with monitoring.MonitoredTimer(
+        counter.get_cell(),
+        monitored_section_name='action1',
+        avoid_repetitive_counting=True,
+    ):
+      time.sleep(1)
+      with monitoring.MonitoredTimer(
+          counter.get_cell(),
+          monitored_section_name='action1',
+          avoid_repetitive_counting=True,
+      ):
+        time.sleep(1)
+
+      # The inner section is not timed.
+      self.assertEqual(counter.get_cell().value(), 0)
+
+    self.assertGreater(counter.get_cell().value(), 0)
+
+  def test_function_decorator(self):
+    counter = monitoring.Counter('test/funcdecorator', 'test func decorator')
+
+    @monitoring.monitored_timer(counter.get_cell())
+    def timed_function(seconds):
+      time.sleep(seconds)
+
+    timed_function(0.001)
+    self.assertGreater(counter.get_cell().value(), 1000)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/ops_test.py

@@ -0,0 +1,515 @@
+# Copyright 2017 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 operations in eager execution."""
+import gc
+import threading
+import weakref
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import execute
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import sparse_ops
+
+
+class OpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  def testExecuteBasic(self):
+    three = constant_op.constant(3)
+    five = constant_op.constant(5)
+    product = three * five
+    self.assertAllEqual(15, product)
+
+  @test_util.run_gpu_only
+  def testMatMulGPU(self):
+    three = constant_op.constant([[3.]]).gpu()
+    five = constant_op.constant([[5.]]).gpu()
+    product = math_ops.matmul(three, five)
+    self.assertEqual([[15.0]], product.numpy())
+
+  def testExecuteStringAttr(self):
+    three = constant_op.constant(3.0)
+    checked_three = array_ops.check_numerics(three,
+                                             message='just checking')
+    self.assertEqual([[3]], checked_three.numpy())
+
+  def testExecuteFloatAttr(self):
+    three = constant_op.constant(3.0)
+    almost_three = constant_op.constant(2.8)
+    almost_equal = math_ops.approximate_equal(
+        three, almost_three, tolerance=0.3)
+    self.assertTrue(almost_equal)
+
+  def testExecuteIntAttr(self):
+    three = constant_op.constant(3)
+    four = constant_op.constant(4)
+    total = math_ops.add_n([three, four])
+    self.assertAllEqual(7, total)
+
+  def testExecuteBoolAttr(self):
+    three = constant_op.constant([[3]])
+    five = constant_op.constant([[5]])
+    product = math_ops.matmul(three, five, transpose_a=True)
+    self.assertAllEqual([[15]], product)
+
+  def testExecuteOneListOutput(self):
+    split_dim = constant_op.constant(1)
+    value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
+    x1, x2, x3 = array_ops.split(value, 3, axis=split_dim)
+    self.assertAllEqual([[0], [3]], x1)
+    self.assertAllEqual([[1], [4]], x2)
+    self.assertAllEqual([[2], [5]], x3)
+
+  def testGraphMode(self):
+    graph = ops.Graph()
+    with graph.as_default(), context.graph_mode():
+      array_ops.placeholder(dtypes.int32)
+    self.assertLen(graph.get_operations(), 1)
+
+  # See comments on handling of int32 tensors on GPU in
+  # EagerTensor.__init__.
+  @test_util.run_gpu_only
+  def testInt32CPUDefault(self):
+    with context.device('/gpu:0'):
+      r = constant_op.constant(1) + constant_op.constant(2)
+    self.assertAllEqual(r, 3)
+
+  def testExecuteListOutputLen1(self):
+    split_dim = constant_op.constant(1)
+    value = constant_op.constant([[0, 1, 2], [3, 4, 5]])
+    result = array_ops.split(value, 1, axis=split_dim)
+    self.assertIsInstance(result, list)
+    self.assertLen(result, 1)
+    self.assertAllEqual([[0, 1, 2], [3, 4, 5]], result[0])
+
+  def testExecuteListOutputLen0(self):
+    empty = constant_op.constant([], dtype=dtypes.int32)
+    result = array_ops_stack.unstack(empty, 0)
+    self.assertIsInstance(result, list)
+    self.assertEmpty(result)
+
+  def testExecuteMultipleNonListOutput(self):
+    x = constant_op.constant([1, 2, 3, 4, 5, 6])
+    y = constant_op.constant([1, 3, 5])
+    result = array_ops.listdiff(x, y)
+    out, idx = result
+    self.assertIs(out, result.out)
+    self.assertIs(idx, result.idx)
+    self.assertAllEqual([2, 4, 6], out)
+    self.assertAllEqual([1, 3, 5], idx)
+
+  def testExecuteMultipleListOutput(self):
+    split_dim = constant_op.constant(1, dtype=dtypes.int64)
+    indices = constant_op.constant([[0, 2], [0, 4], [0, 5], [1, 0], [1, 1]],
+                                   dtype=dtypes.int64)
+    values = constant_op.constant([2, 3, 5, 7, 11])
+    shape = constant_op.constant([2, 7], dtype=dtypes.int64)
+    result = sparse_ops.gen_sparse_ops.sparse_split(
+        split_dim,
+        indices,
+        values,
+        shape,
+        num_split=2)
+    output_indices, output_values, output_shape = result
+    self.assertLen(output_indices, 2)
+    self.assertLen(output_values, 2)
+    self.assertLen(output_shape, 2)
+    self.assertEqual(output_indices, result.output_indices)
+    self.assertEqual(output_values, result.output_values)
+    self.assertEqual(output_shape, result.output_shape)
+    self.assertAllEqual([[0, 2], [1, 0], [1, 1]], output_indices[0])
+    self.assertAllEqual([[0, 0], [0, 1]], output_indices[1])
+    self.assertAllEqual([2, 7, 11], output_values[0])
+    self.assertAllEqual([3, 5], output_values[1])
+    self.assertAllEqual([2, 4], output_shape[0])
+    self.assertAllEqual([2, 3], output_shape[1])
+
+  # TODO(josh11b): Test an op that has multiple outputs, some but not
+  # all of which are lists. Examples: barrier_take_many (currently
+  # unsupported since it uses a type list) or sdca_optimizer (I don't
+  # have an example of legal inputs & outputs).
+
+  def testComposition(self):
+    x = constant_op.constant(1, dtype=dtypes.int32)
+    three_x = x + x + x
+    self.assertEqual(dtypes.int32, three_x.dtype)
+    self.assertAllEqual(3, three_x)
+
+  def testOperatorOverrides(self):
+
+    def ops_test(v1, v2):
+      a = constant_op.constant(v1)
+      b = constant_op.constant(v2)
+
+      self.assertAllEqual((-a), np.negative(v1))
+      self.assertAllEqual(abs(b), np.absolute(v2))
+
+      self.assertAllEqual((a + b), np.add(v1, v2))
+      self.assertAllEqual((a - b), np.subtract(v1, v2))
+      self.assertAllEqual((a * b), np.multiply(v1, v2))
+      self.assertAllEqual((a * a), np.multiply(v1, v1))
+
+      if all(x >= 0 for x in v2):
+        self.assertAllEqual((a**b), np.power(v1, v2))
+      self.assertAllEqual((a / b), np.true_divide(v1, v2))
+
+      self.assertAllEqual((a / a), np.true_divide(v1, v1))
+      self.assertAllEqual((a % b), np.mod(v1, v2))
+
+      self.assertAllEqual((a < b), np.less(v1, v2))
+      self.assertAllEqual((a <= b), np.less_equal(v1, v2))
+      self.assertAllEqual((a > b), np.greater(v1, v2))
+      self.assertAllEqual((a >= b), np.greater_equal(v1, v2))
+
+      # TODO(b/120678848): Remove the else branch once we enable
+      # tensor.Tensor._USE_EQUALITY by default.
+      if tensor.Tensor._USE_EQUALITY:
+        self.assertAllEqual((a == b), np.equal(v1, v2))
+        self.assertAllEqual((a != b), np.not_equal(v1, v2))
+      else:
+        self.assertAllEqual((a == b), np.equal(v1, v2)[0])
+        self.assertAllEqual((a != b), np.not_equal(v1, v2)[0])
+
+      self.assertAllEqual(v1[0], a[constant_op.constant(0)])
+
+    ops_test([1, 4, 8], [2, 3, 5])
+    ops_test([1, -4, -5], [-2, 3, -6])
+
+  def test_basic_slice(self):
+    npt = np.arange(1, 19, dtype=np.float32).reshape(3, 2, 3)
+    t = constant_op.constant(npt)
+
+    self.assertAllEqual(npt[:, :, :], t[:, :, :])
+    self.assertAllEqual(npt[::, ::, ::], t[::, ::, ::])
+    self.assertAllEqual(npt[::1, ::1, ::1], t[::1, ::1, ::1])
+    self.assertAllEqual(npt[::1, ::5, ::2], t[::1, ::5, ::2])
+    self.assertAllEqual(npt[::-1, :, :], t[::-1, :, :])
+    self.assertAllEqual(npt[:, ::-1, :], t[:, ::-1, :])
+    self.assertAllEqual(npt[:, :, ::-1], t[:, :, ::-1])
+    self.assertAllEqual(npt[-2::-1, :, ::1], t[-2::-1, :, ::1])
+    self.assertAllEqual(npt[-2::-1, :, ::2], t[-2::-1, :, ::2])
+
+  def testDegenerateSlices(self):
+    npt = np.arange(1, 19, dtype=np.float32).reshape(3, 2, 3)
+    t = constant_op.constant(npt)
+    # degenerate by offering a forward interval with a negative stride
+    self.assertAllEqual(npt[0:-1:-1, :, :], t[0:-1:-1, :, :])
+    # degenerate with a reverse interval with a positive stride
+    self.assertAllEqual(npt[-1:0, :, :], t[-1:0, :, :])
+    # empty interval in every dimension
+    self.assertAllEqual(npt[-1:0, 2:2, 2:3:-1], t[-1:0, 2:2, 2:3:-1])
+
+  def testEllipsis(self):
+    npt = np.array(
+        [[[[[1, 2], [3, 4], [5, 6]]], [[[7, 8], [9, 10], [11, 12]]]]])
+    t = constant_op.constant(npt)
+
+    self.assertAllEqual(npt[0:], t[0:])
+    # implicit ellipsis
+    self.assertAllEqual(npt[0:, ...], t[0:, ...])
+    # ellipsis alone
+    self.assertAllEqual(npt[...], t[...])
+    # ellipsis at end
+    self.assertAllEqual(npt[0:1, ...], t[0:1, ...])
+    # ellipsis at begin
+    self.assertAllEqual(npt[..., 0:1], t[..., 0:1])
+    # ellipsis at middle
+    self.assertAllEqual(npt[0:1, ..., 0:1], t[0:1, ..., 0:1])
+
+  def testShrink(self):
+    npt = np.array([[[[[1, 2, 4, 5], [5, 6, 7, 8], [9, 10, 11, 12]]],
+                     [[[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]]]])
+    t = constant_op.constant(npt)
+    self.assertAllEqual(npt[:, :, :, :, 3], t[:, :, :, :, 3])
+    self.assertAllEqual(npt[..., 3], t[..., 3])
+    self.assertAllEqual(npt[:, 0], t[:, 0])
+    self.assertAllEqual(npt[:, :, 0], t[:, :, 0])
+
+  @test_util.run_gpu_only
+  def testOpWithInputsOnDifferentDevices(self):
+    # The GPU kernel for the Reshape op requires that the
+    # shape input be on CPU.
+    value = constant_op.constant([1., 2.]).gpu()
+    shape = constant_op.constant([2, 1])
+    reshaped = array_ops.reshape(value, shape)
+    self.assertAllEqual([[1], [2]], reshaped.cpu())
+
+  def testInt64(self):
+    # Fill requires the first input to be an int32 tensor.
+    self.assertAllEqual(
+        [1.0, 1.0],
+        array_ops.fill(constant_op.constant([2], dtype=dtypes.int64),
+                       constant_op.constant(1)))
+
+  @test_util.run_gpu_only
+  def testOutputOnHostMemory(self):
+    # The Shape op kernel on GPU places the output in host memory.
+    value = constant_op.constant([1.]).gpu()
+    shape = array_ops.shape(value)
+    self.assertEqual([1], shape.numpy())
+
+  @test_util.run_gpu_only
+  def testSilentCopy(self):
+    # Temporarily replace the context
+    # pylint: disable=protected-access
+    old_context = context.context()
+    context._set_context(context.Context())
+    try:
+      config.set_device_policy('silent')
+      cpu_tensor = constant_op.constant(1.0)
+      gpu_tensor = cpu_tensor.gpu()
+      self.assertAllEqual(cpu_tensor + gpu_tensor, 2.0)
+    finally:
+      context._set_context(old_context)
+    # pylint: enable=protected-access
+
+  @test_util.run_gpu_only
+  def testSoftPlacement(self):
+    # Temporarily replace the context
+    # pylint: disable=protected-access
+    old_context = context.context()
+    context._set_context(context.Context())
+    try:
+      config.set_device_policy('silent')
+      config.set_soft_device_placement(True)
+      cpu_tensor = constant_op.constant(1.0)
+      result = cpu_tensor + cpu_tensor
+      self.assertEqual(result.device,
+                       '/job:localhost/replica:0/task:0/device:GPU:0')
+    finally:
+      context._set_context(old_context)
+    # pylint: enable=protected-access
+
+  def testRandomUniform(self):
+    scalar_shape = constant_op.constant([], dtype=dtypes.int32)
+
+    x = random_ops.random_uniform(scalar_shape)
+    self.assertEqual(0, x.shape.ndims)
+    self.assertEqual(dtypes.float32, x.dtype)
+
+    x = random_ops.random_uniform(
+        scalar_shape, minval=constant_op.constant(5.),
+        maxval=constant_op.constant(6.))
+    self.assertLess(x, 6)
+    self.assertGreaterEqual(x, 5)
+
+  def testArgsToMatchingEagerDefault(self):
+    # Uses default
+    ctx = context.context()
+    allowed_dtypes = [dtypes.int32, dtypes.int64]
+
+    # Follows standard int conversion rules
+    t, r = execute.args_to_matching_eager([[3, 4]], ctx, allowed_dtypes,
+                                          dtypes.int32)
+    self.assertEqual(t, dtypes.int32)
+    self.assertEqual(r[0].dtype, dtypes.int32)
+    t, r = execute.args_to_matching_eager([[3, 4]], ctx, allowed_dtypes,
+                                          dtypes.int64)
+    self.assertEqual(t, dtypes.int32)
+    self.assertEqual(r[0].dtype, dtypes.int32)
+    # Use int64 since it is a better fit
+    t, r = execute.args_to_matching_eager([[2**48]], ctx, allowed_dtypes,
+                                          dtypes.int32)
+    self.assertEqual(t, dtypes.int64)
+    self.assertEqual(r[0].dtype, dtypes.int64)
+
+    # When the regular tensor conversion fails, then use the default type as a
+    # hint.
+    allowed_dtypes = [dtypes.uint32, dtypes.uint32]
+    t, r = execute.args_to_matching_eager([[3, 4]], ctx, allowed_dtypes,
+                                          dtypes.uint32)
+    self.assertEqual(t, dtypes.uint32)
+    self.assertEqual(r[0].dtype, dtypes.uint32)
+    t, r = execute.args_to_matching_eager([[3, 4]], ctx, allowed_dtypes,
+                                          dtypes.uint64)
+    self.assertEqual(t, dtypes.uint64)
+    self.assertEqual(r[0].dtype, dtypes.uint64)
+
+    t, r = execute.args_to_matching_eager([], ctx, allowed_dtypes, dtypes.int64)
+    self.assertEqual(t, dtypes.int64)
+
+    # Doesn't use default
+    allowed_dtypes = [dtypes.int32, dtypes.string]
+    t, r = execute.args_to_matching_eager([['string', 'arg']], ctx,
+                                          allowed_dtypes, dtypes.int32)
+    self.assertEqual(t, dtypes.string)
+    self.assertEqual(r[0].dtype, dtypes.string)
+
+  def testIdentity(self):
+    self.assertAllEqual(2, array_ops.identity(2))
+
+  @test_util.run_gpu_only
+  def testIdentityOnVariable(self):
+    with context.device('/gpu:0'):
+      v = resource_variable_ops.ResourceVariable(True)
+    self.assertAllEqual(True, array_ops.identity(v))
+
+  def testIncompatibleSetShape(self):
+    x = constant_op.constant(1)
+    with self.assertRaises(ValueError):
+      x.set_shape((1, 2))
+
+  def testCompatibleSetShape(self):
+    x = constant_op.constant([[1, 2]])
+    x.set_shape(tensor_shape.TensorShape([None, 2]))
+    self.assertEqual(x.get_shape(), (1, 2))
+
+  @parameterized.named_parameters(
+      ('Tensor', lambda: constant_op.constant(1.3+1j)),
+      ('Variable', lambda: resource_variable_ops.ResourceVariable(1.3+1j)))
+  def testCastToPrimitiveTypesFrom(self, value_fn):
+    x = value_fn()
+    self.assertIsInstance(int(x), int)
+    self.assertEqual(int(x), 1)
+    self.assertIsInstance(float(x), float)
+    self.assertAllClose(float(x), 1.3)
+    self.assertIsInstance(complex(x), complex)
+    self.assertAllClose(complex(x), 1.3+1j)
+
+  def testCastNonScalarToPrimitiveTypesFails(self):
+    x = constant_op.constant([1.3, 2])
+    with self.assertRaises(TypeError):
+      int(x)
+    with self.assertRaises(TypeError):
+      float(x)
+
+  def testRange(self):
+    x = constant_op.constant(2)
+    self.assertEqual([0, 1], list(range(x)))
+
+  def testFormatString(self):
+    x = constant_op.constant(3.1415)
+    self.assertEqual('3.14', '{:.2f}'.format(x))
+
+  def testNoOpIsNone(self):
+    self.assertIsNone(control_flow_ops.no_op())
+
+  def testEagerContextPreservedAcrossThreads(self):
+    def init_fn():
+      self.assertTrue(context.executing_eagerly())
+      with ops.init_scope():
+        self.assertTrue(context.executing_eagerly())
+        context_switches = context.context().context_switches
+        self.assertLen(context_switches.stack, 1)
+        self.assertFalse(context_switches.stack[0].is_building_function)
+        self.assertEqual(context_switches.stack[0].enter_context_fn,
+                         context.eager_mode)
+
+    self.assertTrue(context.executing_eagerly())
+    t1 = threading.Thread(target=init_fn)
+    t1.start()
+    t1.join()
+
+  def testWeakrefEagerTensor(self):
+    x = constant_op.constant([[1.]])
+    x.at1 = constant_op.constant([[2.]])
+    x.at2 = 3.
+    weak_x = weakref.ref(x)
+    weak_xat1 = weakref.ref(x.at1)
+    del x
+    self.assertIs(weak_x(), None)
+    self.assertIs(weak_xat1(), None)
+
+  def testWeakKeyDictionaryTensor(self):
+    weak_key_dict = weakref.WeakKeyDictionary()
+
+    strong_x = constant_op.constant([[1.]])
+    strong_y = constant_op.constant([[2.]])
+    strong_x_ref = strong_x.ref()
+    strong_y_ref = strong_y.ref()
+    weak_key_dict[strong_x_ref] = constant_op.constant([[3.]])
+    weak_key_dict[strong_y_ref] = constant_op.constant([[4.]])
+    strong_y.a = constant_op.constant([[5.]])
+    weak_x_ref = weakref.ref(strong_x)
+
+    del strong_x, strong_x_ref
+    self.assertIs(weak_x_ref(), None)
+    self.assertEqual([strong_y_ref], list(weak_key_dict))
+    self.assertLen(list(weak_key_dict), 1)
+    self.assertLen(weak_key_dict, 1)
+
+    del strong_y, strong_y_ref
+    self.assertEqual([], list(weak_key_dict))
+
+  def testEagerTensorsCanBeGarbageCollected(self):
+    x = constant_op.constant([[1.]])
+    y = constant_op.constant([[2.]])
+    x.y = y
+    y.x = x
+    weak_x = weakref.ref(x)
+    weak_y = weakref.ref(y)
+    del x
+    del y
+    # Run a gc a few times to ensure cycles are resolved.
+    gc.collect()
+    gc.collect()
+    gc.collect()
+    gc.collect()
+    self.assertIs(weak_x(), None)
+    self.assertIs(weak_y(), None)
+
+  @test_util.disable_tfrt(
+      'b/153697193: tfrt cannot decode python stacktrace yet')
+  # TODO(b/234153596): Disabled because it invalidates stack traces on other
+  # tests (due to partial migration to absl::Status).
+  def DISABLED_testAsyncExceptionStackTrace(self):
+    config.set_synchronous_execution(False)
+
+    def exception_originated_from_here():
+      # Invalid shapes for matmul.
+      return math_ops.matmul([[1]], [[2], [3]])
+
+    # In sync mode, an exception would have been raised here but since this is
+    # in async, the exception will be raised next.
+    x = exception_originated_from_here()
+
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                'in exception_originated_from_here'):
+      x.numpy()
+
+    context.async_clear_error()
+    config.set_synchronous_execution(True)
+
+  def testCrossContextTensorCache(self):
+    old_context = context.context()
+    old_x = constant_op.constant(9.5)
+    context._set_context(context.Context())
+
+    try:
+      new_x = constant_op.constant(9.5)
+      self.assertEqual(new_x.numpy(), 9.5)
+    finally:
+      context._set_context(old_context)
+
+    self.assertEqual(old_x.numpy(), 9.5)
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/polymorphic_function/BUILD

@@ -0,0 +1,638 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "attributes",
+    srcs = ["attributes.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow/python:__subpackages__"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/util:compat",
+    ],
+)
+
+py_library(
+    name = "autograph_util",
+    srcs = ["autograph_util.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/impl:api",
+        "//tensorflow/python/util:tf_decorator_py",
+    ],
+)
+
+py_library(
+    name = "transform",
+    srcs = ["transform.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow/python:__subpackages__"],
+    deps = [],
+)
+
+pytype_strict_library(
+    name = "atomic_function",
+    srcs = ["atomic_function.py"],
+    visibility = ["//tensorflow/python:__subpackages__"],
+    deps = [
+        ":attributes",
+        ":function_type_utils",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/python/client:pywrap_tf_session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:record",
+        "//tensorflow/python/framework:auto_control_deps_utils",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:error_interpolation",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:function_def_to_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:handle_data_util",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:function_utils",
+        "//tensorflow/python/util:tf_stack",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "atomic_function_test",
+    size = "medium",
+    srcs = ["atomic_function_test.py"],
+    deps = [
+        ":atomic_function",
+        ":polymorphic_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "concrete_function",
+    srcs = ["concrete_function.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow/python:__subpackages__"],
+    deps = [
+        ":atomic_function",
+        ":attributes",
+        ":function_type_utils",
+        ":saved_model_exported_concrete",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/python:pywrap_tfe",
+        "//tensorflow/python/eager:backprop_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:forwardprop_util",
+        "//tensorflow/python/eager:graph_only_ops",
+        "//tensorflow/python/eager:record",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:default_gradient",
+        "//tensorflow/python/ops:gradients_util",
+        "//tensorflow/python/ops:handle_data_util",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:trace",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:object_identity",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "concrete_function_test",
+    size = "small",
+    srcs = ["concrete_function_test.py"],
+    deps = [
+        ":atomic_function",
+        ":concrete_function",
+        ":polymorphic_function",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "tracing_compilation",
+    srcs = ["tracing_compilation.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow/python:__subpackages__"],
+    deps = [
+        ":attributes",
+        ":concrete_function",
+        ":function_context",
+        ":function_type_utils",
+        ":tf_method_target",
+        ":transform",
+        "//tensorflow/core/function/capture:capture_container",
+        "//tensorflow/core/function/polymorphism:function_cache",
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/eager:monitoring",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:trace",
+        "//tensorflow/python/util:compat",
+    ],
+)
+
+py_library(
+    name = "tf_method_target",
+    srcs = ["tf_method_target.py"],
+    strict_deps = True,
+    visibility = [
+        "//tensorflow/python/autograph/impl:__pkg__",
+        "//tensorflow/python/eager:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+py_library(
+    name = "polymorphic_function",
+    srcs = ["polymorphic_function.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":attributes",
+        ":autograph_util",
+        ":compiler_ir",
+        ":eager_function_run",
+        ":function_type_utils",
+        ":tf_method_target",
+        ":tracing_compilation",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/function/capture:capture_container",
+        "//tensorflow/core/function/polymorphism:function_cache",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/distribute/parallel_device",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:lift_to_graph",
+        "//tensorflow/python/eager:monitoring",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:control_flow_util",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:trace",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:object_identity",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/python/util:traceback_utils",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "polymorphic_function_test",
+    size = "medium",
+    srcs = ["polymorphic_function_test.py"],
+    shard_count = 15,
+    tags = [
+        "nomac",  # b/157056289
+    ],
+    deps = [
+        ":attributes",
+        ":polymorphic_function",
+        "//tensorflow/core/function/capture:capture_container",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:cancellation",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:lift_to_graph",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:extension_type",
+        "//tensorflow/python/framework:function",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/framework:test_ops",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:clip_ops",
+        "//tensorflow/python/ops:cond_v2",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:data_flow_ops",
+        "//tensorflow/python/ops:functional_ops",
+        "//tensorflow/python/ops:gradients_impl",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:random_ops_gen",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:script_ops",
+        "//tensorflow/python/ops:sendrecv_ops_gen",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops:training_ops_gen",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/ops/structured:structured_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/saved_model:save_context",
+        "//tensorflow/python/saved_model:save_options",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_decorator_py",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "polymorphic_function_test_cpu_only",
+    srcs = ["polymorphic_function_test_cpu_only.py"],
+    # --config=cuda implicitly links in XLA.
+    tags = [
+        "no_cuda_on_cpu_tap",
+        "no_gpu",
+        "no_oss",  # No way to force no XLA linkage in OSS build from here.
+        "no_pip",
+    ],
+    deps = [
+        ":polymorphic_function",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_xla_py_strict_test(
+    name = "polymorphic_function_xla_jit_test",
+    srcs = ["polymorphic_function_xla_jit_test.py"],
+    enabled_backends = [
+        "cpu",
+    ],
+    tags = [
+        "no_mac",
+        "no_oss",  # TODO(b/295654746)
+        "no_pip",
+        "no_windows",
+    ],
+    use_xla_device = False,
+    deps = [
+        ":polymorphic_function",
+        "//tensorflow/compiler/tests:xla_test",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:collective_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:control_flow_util",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops:summary_ops_v2",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+tf_xla_py_strict_test(
+    name = "polymorphic_function_xla_test",
+    srcs = ["polymorphic_function_xla_test.py"],
+    tags = [
+        "no_pip",
+        "no_windows",
+        "nomac",
+    ],
+    deps = [
+        ":polymorphic_function",
+        "//tensorflow/compiler/tests:xla_test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "gradients_test",
+    size = "medium",
+    srcs = ["gradients_test.py"],
+    shard_count = 5,
+    deps = [
+        ":polymorphic_function",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:gradients_impl",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn_grad",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "composite_tensor_utils",
+    srcs = ["composite_tensor_utils.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/util:_pywrap_utils",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+tf_py_strict_test(
+    name = "tracing_compilation_test",
+    size = "medium",
+    srcs = ["tracing_compilation_test.py"],
+    deps = [
+        ":function_type_utils",
+        ":tracing_compilation",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/function/capture:capture_container",
+        "//tensorflow/core/function/polymorphism:function_cache",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/framework:test_ops",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:functional_ops_gen",
+        "//tensorflow/python/ops:gradients_impl",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:resource_variable_ops_gen",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "argument_naming_test",
+    size = "medium",
+    srcs = ["argument_naming_test.py"],
+    deps = [
+        ":polymorphic_function",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "collection_test",
+    size = "medium",
+    srcs = ["collection_test.py"],
+    deps = [
+        ":polymorphic_function",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "eager_function_run",
+    srcs = ["eager_function_run.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "function_context",
+    srcs = ["function_context.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/core/function/polymorphism:function_cache",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:device",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/saved_model:save_context",
+    ],
+)
+
+py_library(
+    name = "saved_model_utils",
+    srcs = ["saved_model_utils.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/saved_model/registration",
+        "//tensorflow/python/trackable:base",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "saved_model_exported_concrete",
+    srcs = ["saved_model_exported_concrete.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":function_type_utils",
+        "//tensorflow/python/trackable:base",
+    ],
+)
+
+py_library(
+    name = "function_type_utils",
+    srcs = ["function_type_utils.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":composite_tensor_utils",
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+        "@six_archive//:six",
+    ],
+)
+
+tf_py_strict_test(
+    name = "function_spec_test",
+    size = "medium",
+    srcs = ["function_spec_test.py"],
+    deps = [
+        ":function_type_utils",
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:tf_decorator_py",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "compiler_ir",
+    srcs = ["compiler_ir.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+tf_xla_py_strict_test(
+    name = "compiler_ir_test",
+    srcs = ["compiler_ir_test.py"],
+    disabled_backends = [
+        "cpu_ondemand",
+        "gpu_a100",
+        "gpu_h100",
+    ],
+    tags = [
+        "no_mac",
+        "no_pip",
+        "no_windows",
+    ],
+    use_xla_device = False,
+    deps = [
+        ":compiler_ir",
+        ":polymorphic_function",
+        "//tensorflow/compiler/tests:xla_test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+    ],
+)

tensorflow/python/eager/polymorphic_function/argument_naming_test.py

@@ -0,0 +1,242 @@
+# Copyright 2017 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class ArgumentNamingTests(test.TestCase, parameterized.TestCase):
+  """Tests for recognizable export signatures from concrete functions."""
+
+  def testBasic(self):
+    @polymorphic_function.function
+    def fn(a, b):
+      return a + b, a * b
+    # Call the function to make def_function happy
+    fn(array_ops.ones([]), array_ops.ones([]))
+
+    fn_op = fn.get_concrete_function(
+        tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32),
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
+    self.assertEqual(
+        ['a', 'b'],
+        [inp.op.name for inp in fn_op.inputs])
+    self.assertEqual(
+        [b'a', b'b'],
+        [inp.op.get_attr('_user_specified_name') for inp in fn_op.inputs])
+    self.assertLen(fn_op.graph.structured_outputs, 2)
+    self.assertAllClose(
+        [3., 2.],
+        fn_op(constant_op.constant(1.), constant_op.constant(2.)))
+    self.assertAllClose(
+        [3., 2.],
+        fn_op(a=constant_op.constant(1.), b=constant_op.constant(2.)))
+
+  def testVariable(self):
+    @polymorphic_function.function
+    def fn(a, b):
+      return a + b, a * b
+    # Call the function to make def_function happy
+    fn(array_ops.ones([]), array_ops.ones([]))
+
+    fn_op = fn.get_concrete_function(
+        tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32),
+        variables.Variable(1.))
+    self.assertEqual(
+        ['a', 'b'],
+        [inp.op.name for inp in fn_op.inputs])
+    self.assertEqual(
+        [b'a', b'b'],
+        [inp.op.get_attr('_user_specified_name') for inp in fn_op.inputs])
+    self.assertLen(fn_op.graph.structured_outputs, 2)
+
+  def testDictReturned(self):
+    @polymorphic_function.function
+    def fn(x, z=(1., 2.), y=3.):
+      z1, z2 = z
+      return {'alpha': x + y + z1, 'beta': x * y + z2}
+    # Call the function to make def_function happy
+    fn(array_ops.ones([]))
+
+    fn_op = fn.get_concrete_function(
+        x=tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32),
+        y=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
+    self.assertEqual(
+        ['x', 'y'],
+        [inp.op.name for inp in fn_op.inputs])
+    self.assertEqual(
+        [b'x', b'y'],
+        [inp.op.get_attr('_user_specified_name') for inp in fn_op.inputs])
+    self.assertEqual({'alpha', 'beta'},
+                     set(fn_op.graph.structured_outputs.keys()))
+
+    fn_op2 = fn.get_concrete_function(
+        z=(tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32,
+                                  name='z_first'),
+           tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32,
+                                  name='z_second')),
+        y=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='custom'),
+        x=4.)
+    self.assertEqual(
+        ['z_first', 'z_second', 'custom'],
+        [inp.op.name for inp in fn_op2.inputs])
+    self.assertEqual(
+        [b'z_first', b'z_second', b'custom'],
+        [inp.op.get_attr('_user_specified_name') for inp in fn_op2.inputs])
+
+    fn_op3 = fn.get_concrete_function(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='custom'),
+        z=(tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.float32,
+                                  name='z1'),
+           tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='z2')),
+        y=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
+    self.assertEqual(
+        ['custom', 'z1', 'z2', 'y'],
+        [inp.op.name for inp in fn_op3.inputs])
+    self.assertEqual(
+        [b'custom', b'z1', b'z2', b'y'],
+        [inp.op.get_attr('_user_specified_name') for inp in fn_op3.inputs])
+
+  def testMethod(self):
+    class HasMethod(object):
+
+      @polymorphic_function.function
+      def method(self, x):
+        return x
+
+    has_method = HasMethod()
+    # Call the function to make def_function happy
+    HasMethod.method(has_method, array_ops.ones([]))
+    class_op = HasMethod.method.get_concrete_function(
+        has_method, tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
+    self.assertEqual(
+        ['x'],
+        [inp.op.name for inp in class_op.inputs])
+    self.assertEqual(
+        [b'x'],
+        [inp.op.get_attr('_user_specified_name') for inp in class_op.inputs])
+    # Call the function to make def_function happy
+    has_method.method(array_ops.ones([]))
+    method_op = has_method.method.get_concrete_function(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32))
+    self.assertEqual(
+        ['x'],
+        [inp.op.name for inp in method_op.inputs])
+    self.assertEqual(
+        [b'x'],
+        [inp.op.get_attr('_user_specified_name') for inp in method_op.inputs])
+    # TODO(allenl): It should be possible to override names when exporting. Do
+    # TensorSpec names need to go in cache keys? Or maybe get_concrete_function
+    # should always retrace?
+    self.skipTest('Not working')
+    method_op = has_method.method.get_concrete_function(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='y'))
+    self.assertEqual(
+        ['y'],
+        [inp.op.name for inp in method_op.inputs])
+    self.assertEqual(
+        [b'y'],
+        [inp.op.get_attr('_user_specified_name') for inp in method_op.inputs])
+
+  def testMethodSignature(self):
+
+    class HasMethod(object):
+
+      @polymorphic_function.function(
+          input_signature=(tensor_spec.TensorSpec(
+              shape=None, dtype=dtypes.float64, name='y'),))
+      def method(self, x):
+        hash(self)  # No weak proxies passed as `self`
+        return x
+
+    has_method = HasMethod()
+    # Call the function to make def_function happy
+    has_method.method(array_ops.ones([], dtype=dtypes.float64))
+    method_op = has_method.method.get_concrete_function()
+    self.assertEqual(
+        ['y'],
+        [inp.op.name for inp in method_op.inputs])
+    self.assertEqual(
+        [b'y'],
+        [inp.op.get_attr('_user_specified_name') for inp in method_op.inputs])
+    method_op2 = has_method.method.get_concrete_function()
+    self.assertEqual(
+        ['y'],
+        [inp.op.name for inp in method_op2.inputs])
+    self.assertEqual(
+        [b'y'],
+        [inp.op.get_attr('_user_specified_name') for inp in method_op2.inputs])
+
+  def testVariadic(self):
+    @polymorphic_function.function
+    def variadic_fn(x, *args, **kwargs):
+      return x + math_ops.add_n(list(args) + list(kwargs.values()))
+
+    # Call the function to make def_function happy
+    variadic_fn(array_ops.ones([]), array_ops.ones([]))
+    variadic_op = variadic_fn.get_concrete_function(
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
+        tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32, name='y'),
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
+        tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32,
+                               name='second_variadic'),
+        z=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
+        zz=tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='cust'))
+    self.assertEqual(
+        ['x', 'y', 'args_1', 'second_variadic', 'z', 'cust'],
+        [inp.op.name for inp in variadic_op.inputs])
+    self.assertEqual(
+        [b'x', b'y', b'args_1', b'second_variadic', b'z', b'cust'],
+        [inp.op.get_attr('_user_specified_name') for inp in variadic_op.inputs])
+
+  def testVariadicInputSignature(self):
+    @polymorphic_function.function(
+        input_signature=(
+            tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32),
+            tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32, name='y'),
+            tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
+            tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name='z'),
+        ))
+    def variadic_fn(x, *args):
+      return x + math_ops.add_n(list(args))
+
+    # Call the function to make def_function happy
+    variadic_fn(array_ops.ones([]), array_ops.ones([]),
+                array_ops.ones([]), array_ops.ones([]))
+    variadic_op = variadic_fn.get_concrete_function()
+    self.assertIn(b'variadic_fn', variadic_op.name)
+    self.assertEqual(
+        ['x', 'y', 'args_1', 'z'],
+        [inp.op.name for inp in variadic_op.inputs])
+    self.assertEqual(
+        [b'x', b'y', b'args_1', b'z'],
+        [inp.op.get_attr('_user_specified_name')
+         for inp in variadic_op.inputs])
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution(
+      config=config_pb2.ConfigProto(device_count={'CPU': 4}))
+  test.main()

tensorflow/python/eager/polymorphic_function/atomic_function.py

@@ -0,0 +1,696 @@
+# 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.
+# ==============================================================================
+"""Implementation for AtomicFunction."""
+
+import dataclasses
+import traceback
+import typing
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.framework import graph_debug_info_pb2
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.python.client import pywrap_tf_session
+from tensorflow.python.eager import context
+from tensorflow.python.eager import record
+from tensorflow.python.eager.polymorphic_function import attributes as attributes_lib
+from tensorflow.python.eager.polymorphic_function import function_type_utils
+from tensorflow.python.framework import auto_control_deps_utils as acd
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import error_interpolation
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import func_graph as func_graph_module
+from tensorflow.python.framework import function_def_to_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import handle_data_util
+from tensorflow.python.types import core
+from tensorflow.python.util import compat
+from tensorflow.python.util import function_utils
+from tensorflow.python.util import tf_stack
+
+
+# TODO(fmuham): Should be lowered to FunctionDef/FunctionRecord.
+@dataclasses.dataclass(frozen=True)
+class CallOptions:
+  """Specifies additional configuration for an AtomicFunction call."""
+
+  # Used by ACD to identify the CollectiveManager this function is scoped in.
+  collective_manager_ids_used: List[int] = dataclasses.field(
+      default_factory=list
+  )
+
+  # Used by ACD to list Ops/Tensors/Callables that must be called in advance.
+  control_captures: List[Any] = dataclasses.field(default_factory=list)
+
+  # Determines what kind of partitioned call is used for this function.
+  is_stateful: bool = False
+
+
+# Maps the (scope_id, name) in runtime to associated AtomicFunctions.
+RUNTIME_FUNCTION_REFS = {}
+
+
+class AtomicFunction(core.AtomicFunction):
+  """A Python callable for functions in the TF Runtime.
+
+  Provides core functionality for tf.function including:
+    - automatic lifecycle management of runtime functions
+    - structured inputs (including captures) and structured outputs
+    - calls from both eager and graph mode
+    - dependency tracking of children functions
+    - runtime error interpolation to identify user code stack traces
+    - control dependencies (including automatic)
+  """
+
+  __slots__ = [
+      "_name",
+      "_bound_context",
+      "_function_type",
+      "_children",
+      "_call_options",
+      "_cached_definition",
+      "_cached_graph",
+      "_generated_graph",
+  ]
+
+  def __init__(
+      self,
+      name: Union[str, bytes],
+      bound_context: context.Context,
+      function_type: function_type_lib.FunctionType,
+      children: Optional[List["AtomicFunction"]] = None,
+      call_options: CallOptions = CallOptions(),
+      cached_graph: Optional[func_graph_module.FuncGraph] = None,
+  ):
+    """Construct a new AtomicFunction.
+
+    Args:
+      name: str/bytes name of the runtime function in the bound context.
+      bound_context: interface to the runtime for the AtomicFunction.
+      function_type: input/output contract for the AtomicFunction
+      children: list of AtomicFunctions that are needed to call this one.
+      call_options: extra configuration options for the call.
+      cached_graph: FuncGraph that this AtomicFunction was generated from (if
+        known). Otherwise it will lazily construct a new corresponding FuncGraph
+        if ever needed.
+    """
+    self._name = compat.as_bytes(name)
+    self._bound_context = bound_context
+    self._function_type = function_type
+    self._children = children if children else []
+    self._call_options = call_options
+    self._cached_definition = None
+
+    self._cached_graph = cached_graph
+    self._generated_graph = None
+
+    ref_key = (self._bound_context.function_scope_id, self.name)
+    if ref_key not in RUNTIME_FUNCTION_REFS:
+      RUNTIME_FUNCTION_REFS[ref_key] = 1
+    else:
+      RUNTIME_FUNCTION_REFS[ref_key] += 1
+
+  @property
+  def name(self) -> bytes:
+    """Name represented in UTF-8 encoded bytes."""
+    return self._name
+
+  @property
+  def function_type(self) -> function_type_lib.FunctionType:
+    """Represents the input/output contract of this function."""
+    return self._function_type
+
+  @property
+  def children(self) -> List["AtomicFunction"]:
+    """AtomicFunctions needed as dependencies for this one."""
+    return self._children
+
+  @property
+  def definition(self) -> function_pb2.FunctionDef:
+    """Current FunctionDef in the Runtime."""
+    return self._bound_context.get_function_def(self.name)
+
+  @property
+  def attributes(self) -> Any:
+    """Returns FunctionDef attributes in the Runtime."""
+    attrs = self.definition.attr
+    # Remove construction context since it is specific to runtime and this fn.
+    attrs.pop(attributes_lib.EAGER_RUNTIME_CONSTRUCTION_CONTEXT, None)
+    return attrs
+
+  @property
+  def graph_debug_info(self) -> graph_debug_info_pb2.GraphDebugInfo:
+    """A GraphDebugInfo proto mapping nodes to corresponding stack traces."""
+    return self._bound_context.get_graph_debug_info(self.name)
+
+  @property
+  def call_options(self) -> CallOptions:
+    """Call options declared for this AtomicFunction."""
+    return self._call_options
+
+  @property
+  def graph_call_attrs(self) -> Dict[str, Any]:
+    """Returns a dictionary of attributes needed to add a call in graph."""
+    attrs = {
+        "is_stateful": self.call_options.is_stateful,
+        "tout": [
+            o.dtype.as_datatype_enum for o in self.function_type.flat_outputs
+        ],
+        "xla_compile_attr": self.cached_definition.attr.get(
+            attributes_lib.XLA_COMPILE, None
+        ),
+    }
+    attrs.update(self._bound_context.function_call_options.as_attrs())
+    return attrs
+
+  @property
+  def _c_func(self) -> Any:
+    """Returns a scoped pybind object containing FunctionRecord in runtime."""
+    return self._bound_context.get_c_function(self.name)
+
+  # TODO(fmuham): Move caching to dependent code and remove method.
+  @property
+  def cached_definition(self) -> function_pb2.FunctionDef:
+    """Cached FunctionDef (not guaranteed to be fresh)."""
+    if self._cached_definition is None:
+      self._cached_definition = self.definition
+
+    return self._cached_definition
+
+  @property
+  def graph(self) -> func_graph_module.FuncGraph:
+    """Returns a FuncGraph corresponding to the AtomicFunction."""
+    if self._cached_graph:
+      return self._cached_graph
+
+    # Lazily generate the graph if one is not specified.
+    if not self._generated_graph:
+      self._generated_graph = to_func_graph(self)
+
+    return self._generated_graph
+
+  def call_with_captures(
+      self, args: Sequence[Any], kwargs: Dict[str, Any], captures: Sequence[Any]
+  ) -> Any:
+    """Calls with args, kwargs, captures and returns structured output."""
+    bound_parameters = self.function_type.bind(*args, **kwargs)
+    tensor_inputs = self.function_type.unpack_inputs(bound_parameters)
+    capture_inputs = self.function_type.unpack_captures(captures)
+    return self.call_preflattened(tensor_inputs + capture_inputs)
+
+  def call_preflattened(self, args: Sequence[core.Tensor]) -> Any:
+    """Calls with flattened tensor inputs and returns the structured output."""
+    flat_outputs = self.call_flat(*args)
+    return self.function_type.pack_output(flat_outputs)
+
+  def call_flat(self, *args: core.Tensor) -> Sequence[core.Tensor]:
+    """Calls with flat tensor inputs and returns flat tensor outputs.
+
+    Args:
+      *args: arguments to call this function with.
+
+    Returns:
+      The outputs of the function call.
+
+    Raises:
+      ValueError: if the number of arguments is incorrect.
+      FunctionAlreadyGarbageCollectedError: if the function is no longer
+        available to be called because it has been garbage collected.
+    """
+    expected_len = len(self.cached_definition.signature.input_arg)
+    if len(args) != expected_len:
+      raise ValueError(
+          f"Signature specifies {expected_len} arguments, got: {len(args)}."
+          f" Expected inputs: {self.cached_definition.signature.input_arg}."
+          f" Received inputs: {args}."
+          f" Function Type: {self.function_type!r}"
+      )
+
+    with InterpolateRuntimeError(self):
+      with ops.control_dependencies(self._call_options.control_captures):
+        # The caller must use record_operation to record this operation in the
+        # eager case, so we enforce the same requirement for the non-eager
+        # case by explicitly pausing recording. We don't have a gradient
+        # registered for PartitionedCall, so recording this operation confuses
+        # forwardprop code (GradientTape manages to ignore it).
+        with record.stop_recording():
+          if self._bound_context.executing_eagerly():
+            outputs = self._bound_context.call_function(
+                self.name,
+                list(args),
+                len(self.function_type.flat_outputs),
+            )
+          else:
+            outputs = make_call_op_in_graph(
+                self,
+                list(args),
+                self._bound_context.function_call_options.as_attrs(),
+            )
+
+    for i, output_type in enumerate(self.function_type.flat_outputs):
+      handle_data = output_type.dtype._handle_data  # pylint: disable=protected-access
+      if handle_data:
+        handle_data_util.set_handle_data(
+            outputs[i], handle_data.shape_inference
+        )
+
+    # TODO(fmuham): Use FunctionType cast here for all cases.
+    if not self._bound_context.executing_eagerly():
+      for i, output_type in enumerate(self.function_type.flat_outputs):
+        outputs[i].set_shape(output_type.shape)
+
+    return outputs
+
+  def __call__(self, *args, **kwargs) -> Any:
+    if self.function_type.captures:
+      raise ValueError(
+          "The FunctionType defines captured inputs. Use call_with_captures"
+          " instead."
+      )
+
+    return self.call_with_captures(args, kwargs, [])
+
+  def __del__(self):
+    if self._generated_graph:
+      func_graph_module.dismantle_func_graph(self._generated_graph)
+
+    if RUNTIME_FUNCTION_REFS is None:
+      return
+
+    key = (self._bound_context.function_scope_id, self.name)
+    RUNTIME_FUNCTION_REFS[key] -= 1
+    if RUNTIME_FUNCTION_REFS[key] < 0:
+      raise RuntimeError(
+          f"AtomicFunction Refcounting for {self.name} is invalid."
+      )
+
+    if RUNTIME_FUNCTION_REFS[key] == 0:
+      try:
+        self._bound_context.remove_function(self.name)
+        RUNTIME_FUNCTION_REFS.pop(key)
+      except TypeError:
+        # Suppress some exceptions, mainly for the case when we're running on
+        # module deletion. Things that can go wrong include the context module
+        # already being unloaded, self._handle._handle_data no longer being
+        # valid, and so on. Printing warnings in these cases is silly
+        # (exceptions raised from __del__ are printed as warnings to stderr).
+        pass  # 'NoneType' object is not callable when the handle has been
+        # partially unloaded.
+      except AttributeError:
+        pass  # 'NoneType' object has no attribute 'eager_mode' when context has
+        # been unloaded. Will catch other module unloads as well.
+
+  def __str__(self):
+    return f"<AtomicFunction> {compat.as_str(self.name)}{self.function_type}"
+
+  def __repr__(self):
+    return (
+        f"AtomicFunction(name={self.name},\n"
+        f"bound_context={self._bound_context},\n"
+        f"function_type={self.function_type!r},\n"
+        f"children={self._children!s},\n"
+        f"call_options={self._call_options},\n"
+        f"cached_graph={self._cached_graph})"
+    )
+
+
+def _set_read_only_resource_inputs_attr(
+    op: ops.Operation, func_graph: func_graph_module.FuncGraph
+):
+  """Sets the list of resource inputs which are read-only.
+
+  This is used by AutomaticControlDependencies.
+
+  Args:
+    op: PartitionedCall Operation.
+    func_graph: FuncGraph.
+  """
+  read_only_indices = acd.get_read_only_resource_input_indices_graph(func_graph)
+  ops.set_int_list_attr(
+      op, acd.READ_ONLY_RESOURCE_INPUTS_ATTR, read_only_indices
+  )
+
+
+def partitioned_call_op(
+    name: str,
+    args: Sequence[core.Tensor],
+    is_stateful: bool,
+    tout: Sequence[Any],
+    config: Any = None,
+    executor_type: Optional[str] = None,
+    xla_compile_attr: Any = None,
+) -> ops.Operation:
+  """Generates a function call op respecting device annotations.
+
+  Args:
+    name: Name of the function to call.
+    args: The arguments of the function, including captured inputs.
+    is_stateful: If the function is stateful.
+    tout: a list containing the output dtypes enums
+    config: (Optional) A `tensorflow::ConfigProto` proto, serialized. If `None`,
+      all optimizations are disabled. Currently only handled for eager defined
+      functions.
+    executor_type: (Optional) A string for the name of the executor to be used
+      in the function call. If not set, or set to an empty string, the default
+      tensorflow executor will be used.
+    xla_compile_attr: (Optional) value of the XLA compilation attribute.
+
+  Returns:
+    Returns the operation.
+  """
+  if config is None:
+    config = function_utils.get_disabled_rewriter_config()
+
+  if executor_type is None:
+    executor_type = ""
+
+  # The generated binding returns an empty list for functions that don't
+  # return any Tensors, hence the need to use `create_op` directly.
+  args = [ops.convert_to_tensor(x) for x in args]
+  tin_attr = attr_value_pb2.AttrValue(
+      list=attr_value_pb2.AttrValue.ListValue(
+          type=[x.dtype.as_datatype_enum for x in args]
+      )
+  )
+  tout_attr = attr_value_pb2.AttrValue(
+      list=attr_value_pb2.AttrValue.ListValue(type=tout)
+  )
+  func_attr = attr_value_pb2.AttrValue(
+      func=attr_value_pb2.NameAttrList(name=name)
+  )
+  executor_type_attr = attr_value_pb2.AttrValue(
+      s=compat.as_bytes(executor_type)
+  )
+
+  # When running in graph mode, the graph and function graphs are optimized
+  # (i.e. run through grappler) per the session options, so we can disable any
+  # eager-specific rewriting.
+  config_proto = attr_value_pb2.AttrValue(s=config)
+
+  op_name = "StatefulPartitionedCall" if is_stateful else "PartitionedCall"
+
+  # Propagate the attribute indicating the need to compile from function to the
+  # call itself.
+  op_attrs = {
+      "Tin": tin_attr,
+      "Tout": tout_attr,
+      "f": func_attr,
+      "config_proto": config_proto,
+      "executor_type": executor_type_attr,
+  }
+  if xla_compile_attr is not None:
+    op_attrs[attributes_lib.XLA_COMPILE] = xla_compile_attr
+
+  op = ops.get_default_graph().create_op(
+      op_name, args, tout, name=op_name, attrs=op_attrs
+  )
+  return op
+
+
+def make_call_op_in_graph(
+    atomic: AtomicFunction,
+    tensor_inputs: Sequence[core.Tensor],
+    context_call_attrs: Dict[str, Any],
+):
+  """Adds an AtomicFunction to graph."""
+  graph = ops.get_default_graph()
+  graph._add_function_recursive(atomic)  # pylint: disable=protected-access
+
+  op = partitioned_call_op(  # pytype: disable=wrong-arg-types  # always-use-property-annotation
+      name=atomic.name,
+      args=tensor_inputs,
+      is_stateful=atomic.call_options.is_stateful,
+      tout=[
+          o.dtype.as_datatype_enum for o in atomic.function_type.flat_outputs
+      ],
+      config=context_call_attrs["config_proto"],
+      executor_type=context_call_attrs["executor_type"],
+      xla_compile_attr=atomic.cached_definition.attr.get(
+          attributes_lib.XLA_COMPILE, None
+      ),
+  )
+  _set_read_only_resource_inputs_attr(op, atomic.graph)
+
+  ops.set_int_list_attr(
+      op,
+      acd.COLLECTIVE_MANAGER_IDS,
+      atomic._call_options.collective_manager_ids_used,  # pylint: disable=protected-access
+  )
+
+  return op.outputs
+
+
+def from_function_def(
+    function_def: function_pb2.FunctionDef,
+    function_type: function_type_lib.FunctionType,
+) -> AtomicFunction:
+  """Create a new AtomicFunction from FunctionDef + FunctionType."""
+  bound_context = context.context()
+  if bound_context.has_function(compat.as_bytes(function_def.signature.name)):
+    raise ValueError("Function already registered in context.")
+
+  bound_context.add_function_def(function_def)
+  return AtomicFunction(
+      function_def.signature.name, bound_context, function_type
+  )
+
+
+def from_func_graph(
+    name: Union[str, bytes],
+    graph: func_graph_module.FuncGraph,
+    attrs: Dict[str, attr_value_pb2.AttrValue],
+    function_type: Optional[function_type_lib.FunctionType] = None,
+    overwrite: bool = False,
+) -> AtomicFunction:
+  """Initializes an AtomicFunction from FuncGraph.
+
+  Args:
+    name: str, the name for the created function.
+    graph: Graph, the graph containing the operations in the function
+    attrs: dict mapping names of attributes to their AttrValue values
+    function_type: known FunctionType to use, otherwise one is derived.
+    overwrite: overwrites function definition in the current context if needed
+
+  Returns:
+    An AtomicFunction instance.
+  """
+  if attrs and attributes_lib.IMPLEMENTS in attrs:
+    # The alternative is to silently drop "implements" tag
+    # but it seems likely it would lead to hard to catch bugs.
+    # Another alternative is to make func_body to preserve the order
+    # of arguments if variables are present. Yet another option
+    # is to automatically replace variables as arguments to functions
+    # to v.read_value() whenever "implements" tag is present
+    # Anytime we annotate existing function we probably want to wrap
+    # it with safe read_value for backward compatibility.
+    has_resource_vars = any(
+        inp.dtype == dtypes.resource for inp in graph.inputs
+    )
+
+    captured_inputs = graph.external_captures + graph.deferred_external_captures
+    assert not any(
+        (has_resource_vars, captured_inputs)
+    ), (
+        'Function {name} has "{attr}={value}" attribute and thus can not '
+        "depend on any tensors outside of its signature or modify variables. "
+        "\n\nNote: variables are always captured and cause function "
+        "re-tracing for every variable called.\n"
+        "  inputs: {inputs}\n  captures: {captured}\n\n"
+        "To pass a variable to such function use  "
+        "use variable.read_value().".format(
+            name=graph.name,
+            attr=attributes_lib.IMPLEMENTS,
+            value=attrs[attributes_lib.IMPLEMENTS],
+            inputs=graph.inputs,
+            captured=captured_inputs,
+        )
+    )
+
+  input_ops = set(arg.op for arg in graph.inputs)
+  operations = [op for op in graph.get_operations() if op not in input_ops]
+
+  graph_output_names = graph._output_names  # pylint: disable=protected-access
+  if graph_output_names is not None and all(
+      ops.tensor_id(t) in graph_output_names for t in graph.outputs
+  ):
+    output_names = [
+        compat.as_bytes(graph_output_names[ops.tensor_id(t)])
+        for t in graph.outputs
+    ]
+    if len(set(output_names)) != len(output_names):
+      # There are duplicate names for some reason, probably an invalid
+      # signature. Revert to auto-naming.
+      output_names = []
+  else:
+    output_names = []
+  with graph._c_graph.get() as c_graph:  # pylint: disable=protected-access
+    fn = pywrap_tf_session.TF_GraphToFunction_wrapper(
+        c_graph,
+        compat.as_str(name),
+        False,
+        [o._c_op for o in operations],  # pylint: disable=protected-access
+        [t._as_tf_output() for t in graph.inputs],  # pylint: disable=protected-access
+        [t._as_tf_output() for t in graph.outputs],  # pylint: disable=protected-access
+        output_names,
+        [o._c_op for o in graph.control_outputs],  # pylint: disable=protected-access
+        [],  # control_output_names
+        None,
+        compat.as_str(""),
+    )
+
+  attrs = attributes_lib.parse_func_attrs(attrs or {})
+  for attr_name, attr_value in attrs.items():
+    serialized = attr_value.SerializeToString()
+    pywrap_tf_session.TF_FunctionSetAttrValueProto(
+        fn, compat.as_str(attr_name), serialized
+    )
+
+  name = compat.as_bytes(name)
+  bound_context = context.context()
+
+  if overwrite and bound_context.has_function(name):
+    bound_context.remove_function(name)
+
+  bound_context.add_c_function(fn)
+  pywrap_tf_session.TF_DeleteFunction(fn)
+
+  call_options = CallOptions(
+      collective_manager_ids_used=getattr(
+          graph, "collective_manager_ids_used", []
+      ),
+      control_captures=graph.function_captures.control,
+      is_stateful=any(op._is_stateful for op in operations),  # pylint: disable=protected-access
+  )
+
+  if not function_type:
+    function_type = function_type_utils.derive_from_graph(graph)
+
+  return AtomicFunction(
+      name,
+      bound_context,
+      function_type,
+      list(graph._functions.values()),  # pylint: disable=protected-access,
+      call_options,
+      cached_graph=graph,
+  )
+
+
+def to_func_graph(atomic: AtomicFunction) -> func_graph_module.FuncGraph:
+  """Generate a FuncGraph from an AtomicFunction."""
+  # pylint: disable=protected-access
+  input_signature, output_signature = function_type_lib.to_structured_signature(
+      atomic.function_type
+  )
+
+  with ops.Graph().as_default():
+    # Insert dependencies in the default graph so the new graph can pull them.
+    for f in atomic.children:
+      ops.get_default_graph()._add_function(f)
+
+    result = function_def_to_graph.function_def_to_graph(
+        atomic.definition,
+        structured_input_signature=input_signature,
+        structured_outputs=output_signature,
+        propagate_device_spec=True,
+        include_library_functions=False,
+    )
+    for f in atomic.children:
+      result._add_function(f)
+
+  # Set input shapes and handle data
+  for i, input_type in enumerate(atomic.function_type.flat_inputs):
+    handle_data = input_type.dtype._handle_data
+    if handle_data:
+      handle_data_util.set_handle_data(
+          result.inputs[i], handle_data.shape_inference
+      )
+    result.inputs[i].set_shape(input_type.shape)
+
+  # Set output shapes and handle data
+  for i, output_type in enumerate(atomic.function_type.flat_outputs):
+    handle_data = output_type.dtype._handle_data
+    if handle_data:
+      handle_data_util.set_handle_data(
+          result.outputs[i], handle_data.shape_inference
+      )
+    result.outputs[i].set_shape(output_type.shape)
+
+  result.collective_manager_ids_used = (
+      atomic.call_options.collective_manager_ids_used,
+  )
+
+  # pylint: enable=protected-access
+  return result
+
+
+class InterpolateRuntimeError(object):
+  """Context Manager that interpolates exceptions received by AtomicFunction."""
+
+  DENY_LIST_PHRASES = ["<embedded"]
+
+  def __init__(self, top_level_func):
+    self._func = top_level_func
+
+  def interpolate(self, message, node_names, graph_debug_info):
+    """Uses the GraphDebugInfo to generate an error message."""
+    error_message = ["Graph execution error:", ""]
+    traces = tf_stack.LoadTracesFromDebugInfo(graph_debug_info)
+
+    for node_name in node_names:
+      error_message.append(
+          f"Detected at node {node_name} defined at (most recent call last):"
+      )
+      if node_name in traces:
+        stack_trace = traces[node_name]
+        for formatted_frame in traceback.format_list(stack_trace):
+          if not any(p in formatted_frame for p in self.DENY_LIST_PHRASES):
+            error_message.append(formatted_frame)
+      else:
+        error_message.append("<stack traces unavailable>")
+
+    error_message.append(message.strip())
+    return "\n".join(error_message)
+
+  def __enter__(self):
+    pass
+
+  def __exit__(self, typ, exc, tb):
+    if not exc or not isinstance(exc, errors.OpError):
+      return False
+    exc = typing.cast(errors.OpError, exc)
+    message = compat.as_text(exc.message)
+    parsed_message, func_tags, node_tags = error_interpolation.parse_message(
+        message
+    )
+    deepest_func = None
+    for func_tag in func_tags:
+      if func_tag.name == compat.as_str(self._func.name):
+        deepest_func = self._func
+      elif deepest_func:
+        next_func = None
+        for child_func in deepest_func.children:
+          if func_tag.name == compat.as_str(child_func.name):
+            next_func = child_func
+            break
+        if next_func is not None and isinstance(next_func, AtomicFunction):
+          deepest_func = next_func
+    if deepest_func:
+      exc._message = self.interpolate(
+          parsed_message,
+          [t.name for t in node_tags],
+          deepest_func.graph_debug_info,
+      )
+    return False

tensorflow/python/eager/polymorphic_function/atomic_function_test.py

@@ -0,0 +1,180 @@
+# 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.
+# ==============================================================================
+
+from tensorflow.python.eager.polymorphic_function import atomic_function
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.platform import test
+
+
+def get_function_def_and_type(foo, inputs):
+  """Traces `foo` generate the FunctionDef and FunctionType."""
+  concrete = polymorphic_function.function(foo).get_concrete_function(*inputs)
+  atomic = concrete._inference_function
+  return atomic.definition, atomic.function_type
+
+
+class AtomicFunctionTest(test.TestCase):
+
+  def test_call_eager(self):
+    definition, func_type = get_function_def_and_type(
+        lambda x, y: x + y, (constant_op.constant(1), constant_op.constant(2))
+    )
+
+    atomic = atomic_function.from_function_def(definition, func_type)
+
+    self.assertRegex(
+        str(atomic),
+        r"<AtomicFunction> .*(x: TensorSpec.*, y: TensorSpec.*) ->"
+        r" TensorSpec.*",
+    )
+    self.assertRegex(
+        repr(atomic).replace("\n", " "),
+        r"AtomicFunction.*name.*bound_context.*function_type.*"
+        r"children.*call_options.*cached_graph.*",
+    )
+
+    self.assertEqual(
+        atomic.call_flat(constant_op.constant(3), constant_op.constant(4))[
+            0
+        ].numpy(),
+        7,
+    )
+
+  def test_call_graph(self):
+    definition, func_type = get_function_def_and_type(
+        lambda x, y: x + y, (constant_op.constant(1), constant_op.constant(2))
+    )
+
+    atomic = atomic_function.from_function_def(definition, func_type)
+
+    @polymorphic_function.function
+    def foo(a, b):
+      return atomic.call_flat(a, b)[0]
+
+    self.assertEqual(
+        foo(constant_op.constant(3), constant_op.constant(4)).numpy(),
+        7,
+    )
+
+  def test_variable_input_eager(self):
+    definition, func_type = get_function_def_and_type(
+        lambda x, y: x + y,
+        (resource_variable_ops.ResourceVariable(1), constant_op.constant(2)),
+    )
+
+    atomic = atomic_function.from_function_def(definition, func_type)
+
+    self.assertEqual(
+        atomic.call_flat(
+            resource_variable_ops.ResourceVariable(3)._handle,
+            constant_op.constant(4),
+        )[0].numpy(),
+        7,
+    )
+
+  def test_variable_input_graph(self):
+    definition, func_type = get_function_def_and_type(
+        lambda x, y: x + y,
+        (resource_variable_ops.ResourceVariable(1), constant_op.constant(2)),
+    )
+
+    atomic = atomic_function.from_function_def(definition, func_type)
+
+    @polymorphic_function.function
+    def foo(a, b):
+      return atomic.call_flat(a, b)[0]
+
+    self.assertEqual(
+        foo(
+            resource_variable_ops.ResourceVariable(3)._handle,
+            constant_op.constant(4),
+        ).numpy(),
+        7,
+    )
+
+  def test_call_with_captures(self):
+    my_capture = constant_op.constant(2)
+
+    @polymorphic_function.function
+    def foo(x):
+      my_dict = {}
+      my_dict["my_tensor"] = x["my_tensor"]
+      my_dict["my_resource"] = x["my_variable"].handle
+      my_dict["my_capture"] = my_capture
+      my_dict["my_ints"] = x["my_ints"]
+      return my_dict
+
+    structured_inputs = {
+        "my_tensor": constant_op.constant(1),
+        "my_variable": resource_variable_ops.ResourceVariable(1),
+        "my_ints": [1, 2, 3],
+    }
+
+    function_def, function_type = get_function_def_and_type(
+        foo, (structured_inputs,)
+    )
+
+    atomic = atomic_function.from_function_def(function_def, function_type)
+
+    with self.assertRaisesRegex(ValueError, "Use call_with_captures instead."):
+      atomic(structured_inputs)
+
+    result = atomic.call_with_captures((structured_inputs,), {}, [my_capture])
+    self.assertEqual(
+        result["my_tensor"].numpy(), structured_inputs["my_tensor"].numpy()
+    )
+    self.assertEqual(result["my_resource"].dtype, dtypes.resource)
+    self.assertEqual(result["my_capture"].numpy(), my_capture.numpy())
+    self.assertEqual(result["my_ints"][0].numpy(), 1)
+    self.assertEqual(result["my_ints"][1].numpy(), 2)
+    self.assertEqual(result["my_ints"][2].numpy(), 3)
+
+  def test_call(self):
+    @polymorphic_function.function
+    def foo(x):
+      my_dict = {}
+      my_dict["my_tensor"] = x["my_tensor"]
+      my_dict["my_resource"] = x["my_variable"].handle
+      my_dict["my_ints"] = x["my_ints"]
+      return my_dict
+
+    structured_inputs = {
+        "my_tensor": constant_op.constant(1),
+        "my_variable": resource_variable_ops.ResourceVariable(1),
+        "my_ints": [1, 2, 3],
+    }
+
+    function_def, function_type = get_function_def_and_type(
+        foo, (structured_inputs,)
+    )
+
+    atomic = atomic_function.from_function_def(function_def, function_type)
+
+    result = atomic(structured_inputs)
+    self.assertEqual(
+        result["my_tensor"].numpy(), structured_inputs["my_tensor"].numpy()
+    )
+    self.assertEqual(result["my_resource"].dtype, dtypes.resource)
+    self.assertEqual(result["my_ints"][0].numpy(), 1)
+    self.assertEqual(result["my_ints"][1].numpy(), 2)
+    self.assertEqual(result["my_ints"][2].numpy(), 3)
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/polymorphic_function/attributes.py

@@ -0,0 +1,184 @@
+# 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.
+# ==============================================================================
+
+"""This file lists FunctionDef attributes and corresponding allowlists."""
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.util import compat
+
+# IMPORTANT: The usage of all the attributes below should be considered tech
+# debt and new additions to this list are discouraged.
+#
+# Historically, attributes have been used as means to pipe extra information
+# down to runtime that is not related to the actual function definition itself.
+#
+# This information is better layered independently and future work is encouraged
+# to pursue that direction instead.
+
+API_IMPLEMENTS = "api_implements"
+API_PREFERRED_DEVICE = "api_preferred_device"
+BACKWARD_FUNCTION = "backward_function_name"
+DISABLE_ACD = "_disable_acd"
+DISABLE_CALL_SHAPE_INFERENCE = "_disable_call_shape_inference"
+DISABLE_SUMMARIES_AT_RUNTIME = "disable_summaries_at_runtime"
+EAGER_RUNTIME_CONSTRUCTION_CONTEXT = "_construction_context"
+FORWARD_FUNCTION = "forward_function_name"
+GO_BACKWARDS = "go_backwards"
+IMPLEMENTS = "_implements"
+INPUT_SHAPES = "_input_shapes"
+INTS_ON_DEVICE = "experimental_ints_on_device"
+NO_INLINE = "_noinline"
+ORIGINAL_FUNCTION_NAME = "_original_func_name"
+OUTPUTS_ON_OP_DEVICE = "_OutputsOnOpDevice"
+QUANTIZED_COMPOSITE_FUNCTION = "tf_quant.composite_function"
+QUANTIZED_OPS = "tf_quant.quantized_ops"
+RUNS_AT_MOST_ONCE = "function_runs_at_most_once"
+RUNTIME_CONSTANT_OPTIMIZATION = "runtime_constant_optimization"
+SHARED_RENDEZVOUS = "shared_rendezvous"
+TF_DATA_FUNCTION = "_tf_data_function"
+TFTRT_ALLOW_BUILD_AT_RUNTIME = "_tftrt_allow_build_at_runtime"
+TFTRT_CONVERT_FUNCTION = "_tftrt_convert_function"
+TFTRT_IS_DYN_OP = "_tftrt_is_dyn_op"
+TFTRT_LOGGER = "_tftrt_trt_logger_name"
+TFTRT_MAX_BATCH_SIZE = "_tftrt_max_batch_size"
+TFTRT_MAX_CACHED_ENGINES = "_tftrt_max_cached_engines"
+TFTRT_MAX_WORKSPACE_SIZE = "_tftrt_max_workspace_size_bytes"
+TFTRT_MIN_SEGMENT_SIZE = "_tftrt_minimum_segment_size"
+TFTRT_PRECISION_MODE = "_tftrt_precision_mode"
+TFTRT_PROFILE_STRATEGY = "_tftrt_profile_strategy"
+TFTRT_USE_CALIBRATION = "_tftrt_use_calibration"
+TFTRT_USE_IMPLICIT_BATCH = "_tftrt_use_implicit_batch"
+TIME_MAJOR = "time_major"
+XLA_COMPILE = "_XlaMustCompile"
+XLA_COMPILE_OPTIONAL = "_XlaCompile"
+XLA_SCOPE = "_XlaScope"
+XLA_SEPARATE_COMPILED_GRADIENTS = "_XlaSeparateCompiledGradients"
+
+POLYMORPHIC_FUNCTION_ALLOWLIST = frozenset({
+    API_IMPLEMENTS,
+    API_PREFERRED_DEVICE,
+    DISABLE_ACD,
+    DISABLE_SUMMARIES_AT_RUNTIME,
+    GO_BACKWARDS,
+    IMPLEMENTS,
+    INTS_ON_DEVICE,
+    NO_INLINE,
+    RUNS_AT_MOST_ONCE,
+    RUNTIME_CONSTANT_OPTIMIZATION,
+    TF_DATA_FUNCTION,
+    TIME_MAJOR,
+    OUTPUTS_ON_OP_DEVICE,
+})
+
+TRACING_COMPILATION_ALLOWLIST = frozenset().union(
+    POLYMORPHIC_FUNCTION_ALLOWLIST,
+    {
+        SHARED_RENDEZVOUS,
+        XLA_COMPILE,
+    },
+)
+
+MONOMORPHIC_FUNCTION_ALLOWLIST = frozenset().union(
+    TRACING_COMPILATION_ALLOWLIST,
+    {
+        BACKWARD_FUNCTION,
+        DISABLE_CALL_SHAPE_INFERENCE,
+        EAGER_RUNTIME_CONSTRUCTION_CONTEXT,
+        FORWARD_FUNCTION,
+        INPUT_SHAPES,
+        ORIGINAL_FUNCTION_NAME,
+        QUANTIZED_COMPOSITE_FUNCTION,
+        QUANTIZED_OPS,
+        TFTRT_ALLOW_BUILD_AT_RUNTIME,
+        TFTRT_CONVERT_FUNCTION,
+        TFTRT_IS_DYN_OP,
+        TFTRT_LOGGER,
+        TFTRT_MAX_BATCH_SIZE,
+        TFTRT_MAX_CACHED_ENGINES,
+        TFTRT_MAX_WORKSPACE_SIZE,
+        TFTRT_MIN_SEGMENT_SIZE,
+        TFTRT_PRECISION_MODE,
+        TFTRT_PROFILE_STRATEGY,
+        TFTRT_USE_CALIBRATION,
+        TFTRT_USE_IMPLICIT_BATCH,
+        XLA_COMPILE_OPTIONAL,
+        XLA_SCOPE,
+        XLA_SEPARATE_COMPILED_GRADIENTS,
+    },
+)
+
+
+def _parse_func_attr_value(key, value):
+  """Converts a python object to an attr_value_pb2.AttrValue object."""
+  if isinstance(value, attr_value_pb2.AttrValue):
+    return value
+  # bool type check has to happen before int since bool is a subclass of int.
+  elif isinstance(value, bool):
+    return attr_value_pb2.AttrValue(b=value)
+  elif isinstance(value, int):
+    return attr_value_pb2.AttrValue(i=value)
+  elif isinstance(value, float):
+    return attr_value_pb2.AttrValue(f=value)
+  elif isinstance(value, (str, bytes)):
+    return attr_value_pb2.AttrValue(s=compat.as_bytes(value))
+  elif isinstance(value, list):
+    list_value = attr_value_pb2.AttrValue.ListValue()
+    for v in value:
+      if isinstance(v, bool):
+        list_value.b.append(v)
+      elif isinstance(v, int):
+        list_value.i.append(v)
+      elif isinstance(v, float):
+        list_value.f.append(v)
+      elif isinstance(v, (str, bytes)):
+        list_value.s.append(compat.as_bytes(v))
+      else:
+        raise ValueError(
+            f"Attributes for {key} must be bool, int, float, or string. "
+            f"Got {type(v)}."
+        )
+    return attr_value_pb2.AttrValue(list=list_value)
+  else:
+    raise ValueError(
+        f"Attribute {key} must be bool, int, float, string, list, or "
+        f"AttrValue. Got {type(value)}."
+    )
+
+
+def parse_func_attrs(attributes, allowlist=None):
+  """Convert the keyword arguments into function_def attributes.
+
+  Currently only support primitive types: bool, int, float and string.
+
+  Args:
+    attributes: the dictionary of attributes.
+    allowlist: set of attribute names allowed.
+  Returns:
+    A dict of attributes where the key is the name of attribute and the value
+      is the AttrValue proto.
+  Raises:
+    ValueError: If the kwargs contains unallowlisted name or unsupported value
+      types.
+  """
+  if not allowlist:
+    allowlist = MONOMORPHIC_FUNCTION_ALLOWLIST
+
+  attrs = {}
+  for key, value in attributes.items():
+    if key not in allowlist:
+      raise ValueError(
+          f"Allowlist does not support `{key}` as an attribute.")
+    attrs[key] = _parse_func_attr_value(key, value)
+  return attrs

tensorflow/python/eager/polymorphic_function/autograph_util.py

@@ -0,0 +1,59 @@
+# 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.
+# ==============================================================================
+# pylint: disable=unidiomatic-typecheck
+"""Autograph utility functions for polymorphic_function."""
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.impl import api
+from tensorflow.python.util import tf_decorator
+
+
+def py_func_from_autograph(
+    python_func,
+    autograph_options=None,
+):
+  """Compile a python function using autograph, for use with FuncGraph.
+
+  Args:
+    python_func: the Python function to compile.
+    autograph_options: additional knobs to control when `autograph=True`.
+      See https://www.tensorflow.org/guide/autograph for more information.
+  Returns:
+    python_func, converted using autograph.
+  """
+  _, original_func = tf_decorator.unwrap(python_func)
+
+  def autograph_handler(*args, **kwargs):
+    """Calls a converted version of original_func."""
+    try:
+      return api.converted_call(
+          original_func,
+          args,
+          kwargs,
+          options=converter.ConversionOptions(
+              recursive=True,
+              optional_features=autograph_options,
+              user_requested=True,
+          ))
+    except Exception as e:  # pylint:disable=broad-except
+      if hasattr(e, "ag_error_metadata"):
+        raise e.ag_error_metadata.to_exception(e)
+      else:
+        raise
+
+  # Wrapping around a decorator allows checks like tf_inspect.getargspec
+  # to be accurate.
+  converted_func = tf_decorator.make_decorator(original_func, autograph_handler)
+  return tf_decorator.rewrap(python_func, original_func, converted_func)

tensorflow/python/eager/polymorphic_function/collection_test.py

@@ -0,0 +1,69 @@
+# Copyright 2017 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.
+# ==============================================================================
+
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class FunctionCollectionTest(test.TestCase):
+
+  def testCollectionValueAccess(self):
+    """Read values from graph collections inside of defun."""
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g):
+        x = 2
+        y = 5
+        ops.add_to_collection('x', x)
+        ops.add_to_collection('y', y)
+
+        @polymorphic_function.function
+        def fn():
+          x_const = constant_op.constant(ops.get_collection('x')[0])
+          y_const = constant_op.constant(ops.get_collection('y')[0])
+          z = math_ops.add(x_const, y_const)
+          ops.add_to_collection('z', 7)
+          return z
+
+        self.assertEqual(7, int(self.evaluate(fn())))
+        self.assertEqual(ops.get_collection('x'), [2])
+        self.assertEqual(ops.get_collection('y'), [5])
+        self.assertEqual(ops.get_collection('z'), [])
+
+  def testCollectionVariableValueAccess(self):
+    """Read variable value from graph collections inside of defun."""
+    with ops.Graph().as_default() as g:
+      with self.session(graph=g):
+        v = resource_variable_ops.ResourceVariable(1.0)
+
+        @polymorphic_function.function
+        def f():
+          return v.read_value()
+
+        self.evaluate(variables.global_variables_initializer())
+        self.assertEqual(1.0, float(self.evaluate(f())))
+        self.assertLen(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES), 1)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution(
+      config=config_pb2.ConfigProto(device_count={'CPU': 4}))
+  test.main()

tensorflow/python/eager/polymorphic_function/compiler_ir.py

@@ -0,0 +1,130 @@
+# 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.
+# ==============================================================================
+"""Implementation for defining get_compiler_ir."""
+from typing import List, Optional
+import warnings
+
+from tensorflow.core.function import trace_type
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import random_ops
+from tensorflow.python.util import nest
+
+
+def maybe_get_device_name(device_name):
+  # TODO(cheshire): This is a hack to get the current "preferred" device,
+  # there is no current API to get it otherwise.
+  if device_name is None:
+    device_name = random_ops.random_normal([]).device
+  return device_name
+
+
+# TODO(fmuham): Use trace_type._flatten here instead when available
+def make_handledata_tensor_specs(resource_vars):
+  """Convert tf.Variable list to its corresponding TensorSpec list."""
+  if not all(x.dtype is dtypes.resource for x in resource_vars):
+    raise RuntimeError("Resource_vars must be tf.resource list.")
+  inner_context = trace_type.InternalTracingContext()
+  trace_type_inputs = trace_type.from_value(
+      tuple(resource_vars), inner_context
+  ).components
+
+  def to_resource_spec(traced_input):
+    try:
+      handle_data = traced_input.dtype._handle_data.shape_inference  # pylint: disable=protected-access
+      shape_and_type = handle_data.shape_and_type[0]
+      spec = tensor_spec.TensorSpec(
+          shape=shape_and_type.shape, dtype=shape_and_type.dtype
+      )
+      return spec
+    except Exception as e:
+      raise ValueError(
+          "Fail to convert tf.Variable list to TensorSpec list. The error"
+          " is: %s" % e
+      ) from e
+
+  return [to_resource_spec(trace_type) for trace_type in trace_type_inputs]
+
+
+def from_concrete_function(
+    concrete_fn,
+    specialized_flat_specs: Optional[List[tensor_spec.TensorSpec]] = None,
+):
+  """Generate the Compiler Ir from tf concrete function with TensorSpec.
+
+  Args:
+    concrete_fn: returned by using get_concrete_function.
+    specialized_flat_specs: specialized flat tf.TensorSpecs for function args.
+
+  Returns:
+    Function callable that generate the HLO text.
+
+  Raises:
+      ValueError: if concrete_fn is not "compilable" without concrete
+      inputs.
+  """
+  context.ensure_initialized()
+  fn_name = concrete_fn.name
+  filtered_flat_specs = specialized_flat_specs or list(
+      nest.flatten(concrete_fn.structured_input_signature)
+  )
+
+  if not all(s.shape.is_fully_defined() for s in filtered_flat_specs):
+    raise ValueError(
+        f"Only support static input shape but got inputs = {concrete_fn.inputs}"
+    )
+
+  def compiler_ir_generator(stage="hlo", device_name=None, platform_name=None):
+    """Gets the compiler IR bytes.
+
+    Args:
+      stage: The exported stage for the given function.
+      device_name: The name of the device with the form as
+        "/job:localhost/replica:0/task:0/device:CPU:0", "/device:TPU:0" etc.
+        When this is used, actual device is needed for getting the compiler IR.
+      platform_name: The name of the platform, e.g. "TPU". See the comment in
+        `get_compiler_ir` in `context.py`.
+
+    Returns:
+      The compiler IR bytes.
+    """
+    if device_name is not None:
+      if platform_name is not None:
+        raise ValueError(
+            "device_name and platform_name cannot be provided at the same time."
+        )
+      warnings.warn("device_name is being deprecated. Use platform_name.")
+    device_name = maybe_get_device_name(device_name)
+    res_bytes = context.context().get_compiler_ir(
+        device_name=device_name,
+        platform_name=platform_name,
+        function_name=fn_name,
+        flat_args=filtered_flat_specs,
+        captured_inputs=concrete_fn.captured_inputs,
+        stage=stage,
+    )
+    if stage in (
+        # Ordered by IrExportStage enum order
+        "stablehlo_serialized",
+        "hlo_serialized",
+        "optimized_hlo_serialized",
+        "optimized_hlo_proto_serialized",
+    ):
+      return res_bytes
+    else:
+      return res_bytes.decode("utf-8")
+
+  return compiler_ir_generator

tensorflow/python/eager/polymorphic_function/compiler_ir_test.py

@@ -0,0 +1,241 @@
+# 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.
+# ==============================================================================
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager.polymorphic_function import compiler_ir
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.util import nest
+
+
+class CompilerIrTest(xla_test.XLATestCase):
+
+  def _compareTwoMethodsCompilerIROutput(self, f, args, kwargs):
+    flat_args = list(args) + list(kwargs.values())
+    if not all([isinstance(x, tensor.Tensor) for x in flat_args]):
+      self.skipTest('It only support args and kwargs are all tf.Tensor types.')
+
+    args_spec = nest.map_structure(tensor.TensorSpec.from_tensor, args)
+    kwargs_spec = nest.map_structure(tensor.TensorSpec.from_tensor, kwargs)
+
+    hlo_1 = f.experimental_get_compiler_ir(*args, **kwargs)()
+    hlo_2 = f.experimental_get_compiler_ir(*args_spec, **kwargs_spec)()
+
+    if hlo_1 != hlo_2:
+      self.fail(
+          'The tensor_spec way experimental_get_compiler_ir give diff result'
+          ' to normal experimental_get_compiler_ir.'
+          f' \nhlo(concrete_input):\n{hlo_1}\nhlo(tensor_spec):\n{hlo_2}\n'
+      )
+
+    # Check that StableHLO conversion succeeds
+    hlo_3 = f.experimental_get_compiler_ir(*args, **kwargs)(stage='stablehlo')
+    self.assertIn('stablehlo', hlo_3)
+
+    # Check that StableHLO bytecode conversion succeeds.
+    # MLIR bytecode files all begin with magic `MLiR` byte, check for byte.
+    hlo_4 = f.experimental_get_compiler_ir(*args, **kwargs)(
+        stage='stablehlo_serialized'
+    )
+    self.assertIn(b'ML\xefR', hlo_4)
+
+  def test_zero_input(self):
+    with ops.device('device:{}:0'.format(self.device)):
+
+      @polymorphic_function.function(jit_compile=True, autograph=False)
+      def fun_tf():
+        return array_ops.zeros((10), dtype=dtypes.int32)
+
+      self._compareTwoMethodsCompilerIROutput(fun_tf, [], {})
+
+  def test_constant_slice(self):
+    with ops.device('device:{}:0'.format(self.device)):
+      # Constant slice. This is the common case.
+      x = array_ops.zeros((10,), dtype=dtypes.int32)
+
+      @polymorphic_function.function(jit_compile=True, autograph=False)
+      def fun_tf(x):
+        begin = 0
+        return x[begin:5]
+
+      self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
+
+  def test_compile_time_constant(self):
+    with ops.device('device:{}:0'.format(self.device)):
+      # Non-constant slice, but compile-time constant depending only on shapes.
+      x = array_ops.zeros((10,), dtype=dtypes.int32)
+
+      @polymorphic_function.function(jit_compile=True, autograph=False)
+      def fun_tf(x):
+        # begin is a compile-time constant, even if x is not
+        begin = array_ops.shape_v2(x)[0] - 2
+        return x[begin:]
+
+      self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
+
+  def test_capture_constant(self):
+    with ops.device('device:{}:0'.format(self.device)):
+      # Capture a constant
+      outer_ct = [3.0]
+      x = ops.convert_to_tensor([2.0, 3.0, 4.0], dtype=dtypes.float32)
+
+      @polymorphic_function.function(jit_compile=True, autograph=False)
+      def fun_tf(x):
+        return x * gen_array_ops.broadcast_to(outer_ct, x.shape) + 1.0
+
+      self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
+
+  def test_unsupported_dynamic_input(self):
+    with ops.device('device:{}:0'.format(self.device)):
+
+      @polymorphic_function.function(jit_compile=True)
+      def f(x):
+        return x
+
+      with self.assertRaisesRegex(
+          ValueError, 'Only support static input shape but got'
+      ):
+        args_spec = [tensor.TensorSpec((None), dtype=dtypes.float32)]
+        concrete_fn = f.get_concrete_function(*args_spec)
+        _ = compiler_ir.from_concrete_function(concrete_fn)(stage='hlo')
+
+  def test_unsupported_shape_depend_input(self):
+    with ops.device('device:{}:0'.format(self.device)):
+      # Those cases output shapes are dynamic.
+      @polymorphic_function.function(jit_compile=True)
+      def f2(x):
+        return x[x[0] : 0]
+
+      args = [ops.convert_to_tensor([1, 2, 3, 4])]
+      args_spec = nest.map_structure(tensor.TensorSpec.from_tensor, args)
+      concrete_fn = f2.get_concrete_function(*args_spec)
+      _ = compiler_ir.from_concrete_function(concrete_fn)(stage='hlo')
+
+  def test_make_handledata_tensor_specs(self):
+    with ops.device('device:{}:0'.format(self.device)):
+      v1 = variables.Variable([0.1, 0.1])
+      v3 = variables.Variable([1], dtype=dtypes.int32)
+
+      @polymorphic_function.function(jit_compile=True)
+      def f4(a, b):
+        return (a + b) * v1 - math_ops.cast(v3, dtypes.float32)
+
+      a = constant_op.constant([1.1, 1.1])
+      b = constant_op.constant([2.2, 2.2])
+
+      kwargs = {'b': a, 'a': b}
+
+      kwargs_spec = nest.map_structure(
+          tensor.TensorSpec.from_tensor, kwargs
+      )
+      concrete_fn = f4.get_concrete_function(**kwargs_spec)
+      captured_inputs = concrete_fn.captured_inputs
+      captured_spec = compiler_ir.make_handledata_tensor_specs(captured_inputs)
+      self.assertEqual(len(captured_spec), 2)
+      self.assertEqual(
+          captured_spec[0], tensor.TensorSpec((2), dtype=dtypes.float32)
+      )
+      self.assertEqual(
+          captured_spec[1], tensor.TensorSpec((1), dtype=dtypes.int32)
+      )
+
+  def test_capture_variable_1(self):
+    if 'gpu' in self.device.lower():
+      self.skipTest('Skip test on GPU')
+
+    with ops.device('device:{}:0'.format(self.device)):
+      v1 = variables.Variable([0.1, 0.1])
+      v3 = variables.Variable([1], dtype=dtypes.int32)
+
+      @polymorphic_function.function(jit_compile=True)
+      def f4(a, b):
+        return (a + b) * v1 - math_ops.cast(v3, dtypes.float32)
+
+      a = constant_op.constant([1.1, 1.1])
+      b = constant_op.constant([2.2, 2.2])
+
+      kwargs = {'b': a, 'a': b}
+      self._compareTwoMethodsCompilerIROutput(f4, [], kwargs)
+
+  def test_capture_variable_2(self):
+    if 'gpu' in self.device.lower():
+      self.skipTest('Skip test on GPU')
+
+    with ops.device('device:{}:0'.format(self.device)):
+      v2 = variables.Variable(2.0, dtype=dtypes.float32)
+      v3 = variables.Variable(3.0, dtype=dtypes.float32)
+
+      @polymorphic_function.function(jit_compile=True)
+      def fun_tf(x):
+        # Defining tf.constants inside func_body is okay.
+        t4 = constant_op.constant(4.0, dtype=dtypes.float32)
+        t5 = constant_op.constant(5.0, dtype=dtypes.float32)
+        return (x * v3 + t4 + v2) * v3 + t5
+
+      x = constant_op.constant(2.0, dtype=dtypes.float32)
+      self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
+
+  def test_capture_constants(self):
+    if 'gpu' in self.device.lower():
+      self.skipTest('Skip test on GPU')
+
+    with ops.device('device:{}:0'.format(self.device)):
+      v2 = variables.Variable(2.0, dtype=dtypes.float32)
+      v3 = variables.Variable(3.0, dtype=dtypes.float32)
+      t4 = constant_op.constant([4.0, 5.0], dtype=dtypes.float32)
+      t5 = constant_op.constant([5.0, 6.0], dtype=dtypes.float32)
+
+      @polymorphic_function.function(jit_compile=True)
+      def fun_tf(x):
+        return (x * v3 + t4 + v2) * v3 + t5
+
+      x = constant_op.constant([2.0, 3.0], dtype=dtypes.float32)
+      self._compareTwoMethodsCompilerIROutput(fun_tf, [x], {})
+
+  def test_from_concrete_function_with_args(self):
+    with ops.device('device:{}:0'.format(self.device)):
+      v2 = variables.Variable(2.0, dtype=dtypes.float32)
+      v3 = variables.Variable(3.0, dtype=dtypes.float32)
+      # Capturing tf.constants outside func_body is not okay.
+      t4 = constant_op.constant(4.0, dtype=dtypes.float32)
+      t5 = constant_op.constant(5.0, dtype=dtypes.float32)
+
+      @polymorphic_function.function(jit_compile=True)
+      def fun_tf(x):
+        return (x * v3 + t4 + v2) * v3 + t5
+
+      concrete_fn = fun_tf.get_concrete_function(
+          tensor.TensorSpec((None,), dtype=dtypes.float32)
+      )
+
+      x = tensor.TensorSpec((10,), dtype=dtypes.float32)
+      hlo_1 = compiler_ir.from_concrete_function(concrete_fn, [x])(stage='hlo')
+      self.assertIn('f32[10]', hlo_1)
+      x = tensor.TensorSpec((20,), dtype=dtypes.float32)
+      hlo_2 = compiler_ir.from_concrete_function(concrete_fn, [x])(stage='hlo')
+      self.assertIn('f32[20]', hlo_2)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/polymorphic_function/composite_tensor_utils.py

@@ -0,0 +1,39 @@
+# 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.
+# ==============================================================================
+"""Utility to manipulate CompositeTensors in tf.function."""
+
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.util import _pywrap_utils
+from tensorflow.python.util import nest
+
+
+# TODO(b/240337581, b/240337099): Remove this function when we de-alias
+# dt_resource tensors or tf.nest support is_leaf.
+def flatten_with_variables(inputs):
+  """Flattens `inputs` but don't expand `ResourceVariable`s."""
+  # We assume that any CompositeTensors have already converted their components
+  # from numpy arrays to Tensors, so we don't need to expand composites here for
+  # the numpy array conversion. Instead, we do so because the flattened inputs
+  # are eventually passed to ConcreteFunction()._call_flat, which requires
+  # expanded composites.
+  flat_inputs = []
+  for value in nest.flatten(inputs):
+    if (isinstance(value, composite_tensor.CompositeTensor) and
+        not _pywrap_utils.IsResourceVariable(value)):
+      components = value._type_spec._to_components(value)  # pylint: disable=protected-access
+      flat_inputs.extend(flatten_with_variables(components))
+    else:
+      flat_inputs.append(value)
+  return flat_inputs

tensorflow/python/eager/polymorphic_function/concrete_function_test.py

@@ -0,0 +1,157 @@
+# 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.eager.polymorphic_function import atomic_function
+from tensorflow.python.eager.polymorphic_function import concrete_function as cf
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import func_graph as func_graph_module
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class ConcreteFunctionTest(test.TestCase, parameterized.TestCase):
+
+  def concrete_function_with_attrs(self, attrs):
+    func_graph = func_graph_module.FuncGraph("f")
+    return cf.ConcreteFunction.from_func_graph(func_graph, None, attrs=attrs)
+
+  @parameterized.parameters(
+      ({"api_implements": True}, attr_value_pb2.AttrValue(b=True)),
+      ({"api_implements": 1}, attr_value_pb2.AttrValue(i=1)),
+      ({"api_implements": 1.0}, attr_value_pb2.AttrValue(f=1.0)),
+      (
+          {"api_implements": "test"},
+          attr_value_pb2.AttrValue(s=compat.as_bytes("test")),
+      ),
+  )
+  def test_parses_func_attr_scalar_values(self, attrs, expected):
+    self.assertEqual(
+        self.concrete_function_with_attrs(attrs=attrs).function_def.attr[
+            "api_implements"
+        ],
+        expected,
+    )
+
+  def test_parses_func_attr_list_values(self):
+    self.assertProtoEquals(
+        r"""
+        list {
+            s: 'test'
+            b: True
+            i: 1
+            f: 1.0
+        }
+        """,
+        self.concrete_function_with_attrs(
+            attrs={"api_implements": ["test", True, 1, 1.0]}
+        ).function_def.attr["api_implements"],
+    )
+
+  def test_raises_value_error_for_invalid_attr(self):
+    with self.assertRaisesRegex(ValueError, "Attribute api_implements must be"):
+      self.concrete_function_with_attrs(attrs={"api_implements": None})
+
+  def test_generate_from_atomic(self):
+    @polymorphic_function.function
+    def add_dicts(dict_a, dict_b):
+      result = {}
+      for key in dict_a.keys():
+        result[key] = dict_a[key] + dict_b[key]
+      return result
+
+    dict_a = {
+        "tensor": constant_op.constant(1),
+        "variable": variables.Variable(2),
+        "ragged_tensor": ragged_tensor.RaggedTensor.from_row_splits(
+            values=[3, 1, 4, 1, 5, 9, 2, 6], row_splits=[0, 4, 4, 7, 8, 8]
+        ),
+        "python_int": 4,
+    }
+    dict_b = {
+        "tensor": constant_op.constant(2),
+        "variable": variables.Variable(5),
+        "ragged_tensor": ragged_tensor.RaggedTensor.from_row_splits(
+            values=[4, 2, 4, 1, 6, 9, 3, 6], row_splits=[0, 4, 4, 7, 8, 8]
+        ),
+        "python_int": 5,
+    }
+
+    original_concrete_fn = add_dicts.get_concrete_function(dict_a, dict_b)
+
+    # Get the atomic function and delete everything else.
+    atomic_fn = original_concrete_fn._inference_function
+    del add_dicts
+    del original_concrete_fn
+
+    # Regenerate the ConcreteFunction.
+    concrete_fn = cf.ConcreteFunction(atomic_fn)
+    result = concrete_fn(dict_a, dict_b)
+
+    # Call and check results.
+    self.assertEqual(result["tensor"].numpy(), 3)
+    self.assertEqual(result["variable"].numpy(), 7)
+    self.assertEqual(
+        result["ragged_tensor"].flat_values.numpy().tolist(),
+        [7, 3, 8, 2, 11, 18, 5, 12],
+    )
+    self.assertEqual(result["python_int"].numpy(), 9)
+
+  def test_generate_from_def(self):
+    @polymorphic_function.function
+    def add_dicts(dict_a, dict_b):
+      result = {}
+      for key in dict_a.keys():
+        result[key] = dict_a[key] + dict_b[key]
+      return result
+
+    dict_a = {
+        "tensor": constant_op.constant(1),
+        "variable": variables.Variable(2),
+        "python_int": 4,
+    }
+    dict_b = {
+        "tensor": constant_op.constant(2),
+        "variable": variables.Variable(5),
+        "python_int": 5,
+    }
+
+    original_concrete_fn = add_dicts.get_concrete_function(dict_a, dict_b)
+
+    # Get FunctionDef + FunctionType and delete everything else.
+    function_def = original_concrete_fn.function_def
+    function_type = original_concrete_fn.function_type
+    del add_dicts
+    del original_concrete_fn
+
+    # Regenerate the ConcreteFunction.
+    atomic_fn = atomic_function.from_function_def(function_def, function_type)
+    concrete_fn = cf.ConcreteFunction(atomic_fn)
+    result = concrete_fn(dict_a, dict_b)
+
+    # Call and check results.
+    self.assertEqual(result["tensor"].numpy(), 3)
+    self.assertEqual(result["variable"].numpy(), 7)
+    self.assertEqual(result["python_int"].numpy(), 9)
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/polymorphic_function/eager_function_run.py

@@ -0,0 +1,111 @@
+# 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.
+# ==============================================================================
+# pylint: disable=unidiomatic-typecheck
+"""Eager semantics for polymorphic function."""
+
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+RUN_FUNCTIONS_EAGERLY = False
+
+
+@tf_export("config.functions_run_eagerly")
+def functions_run_eagerly():
+  """Returns the value of the `run_functions_eagerly` setting."""
+  return RUN_FUNCTIONS_EAGERLY
+
+
+@tf_export("config.run_functions_eagerly")
+def run_functions_eagerly(run_eagerly):
+  """Enables / disables eager execution of `tf.function`s.
+
+  Calling `tf.config.run_functions_eagerly(True)` will make all
+  invocations of `tf.function` run eagerly instead of running as a traced graph
+  function. This can be useful for debugging. As the code now runs line-by-line,
+  you can add arbitrary `print` messages or pdb breakpoints to monitor the
+  inputs/outputs of each Tensorflow operation. However, you should avoid using
+  this for actual production because it significantly slows down execution.
+
+  >>> def my_func(a):
+  ...  print(f'a: {a}')
+  ...  return a + a
+  >>> a_fn = tf.function(my_func)
+
+  >>> # A side effect the first time the function is traced
+  >>> # In tracing time, `a` is printed with shape and dtype only
+  >>> a_fn(tf.constant(1))
+  a: Tensor("a:0", shape=(), dtype=int32)
+  <tf.Tensor: shape=(), dtype=int32, numpy=2>
+
+  >>> # `print` is a python side effect, it won't execute as the traced function
+  >>> # is called
+  >>> a_fn(tf.constant(2))
+  <tf.Tensor: shape=(), dtype=int32, numpy=4>
+
+  >>> # Now, switch to eager running
+  >>> tf.config.run_functions_eagerly(True)
+  >>> # The code now runs eagerly and the actual value of `a` is printed
+  >>> a_fn(tf.constant(2))
+  a: 2
+  <tf.Tensor: shape=(), dtype=int32, numpy=4>
+
+  >>> # Turn this back off
+  >>> tf.config.run_functions_eagerly(False)
+
+  Note: This flag has no effect on functions passed into tf.data transformations
+  as arguments. tf.data functions are never executed eagerly and are always
+  executed as a compiled Tensorflow Graph.
+
+  Args:
+    run_eagerly: Boolean. Whether to run functions eagerly.
+  """
+  global RUN_FUNCTIONS_EAGERLY
+  RUN_FUNCTIONS_EAGERLY = bool(run_eagerly)
+
+
+@deprecation.deprecated(
+    None, "Use `tf.config.run_functions_eagerly` instead of the experimental "
+    "version.")
+@tf_export("config.experimental_run_functions_eagerly")
+def experimental_run_functions_eagerly(run_eagerly):
+  """Enables / disables eager execution of `tf.function`s.
+
+  Calling `tf.config.experimental_run_functions_eagerly(True)` will make all
+  invocations of `tf.function` run eagerly instead of running as a traced graph
+  function.
+
+  See `tf.config.run_functions_eagerly` for an example.
+
+  Note: This flag has no effect on functions passed into tf.data transformations
+  as arguments. tf.data functions are never executed eagerly and are always
+  executed as a compiled Tensorflow Graph.
+
+  Args:
+    run_eagerly: Boolean. Whether to run functions eagerly.
+
+  Returns:
+    None
+  """
+  return run_functions_eagerly(run_eagerly)
+
+
+@deprecation.deprecated(
+    None,
+    "Use tf.config.functions_run_eagerly instead of the experimental version.")
+@tf_export("config.experimental_functions_run_eagerly")
+def experimental_functions_run_eagerly():
+  """Returns the value of the `experimental_run_functions_eagerly` setting."""
+  return functions_run_eagerly()

tensorflow/python/eager/polymorphic_function/function_context.py

@@ -0,0 +1,127 @@
+# 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.
+# ==============================================================================
+"""Context information for a tf.function."""
+
+from typing import NamedTuple, Any
+
+from tensorflow.core.function.polymorphism import function_cache
+from tensorflow.python.eager import context
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.framework import func_graph as func_graph_module
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.saved_model import save_context
+
+
+# EagerContext is used by tf.function to identify cases where tracing
+# needs to occur due to a change in conditions other than the arguments.
+class EagerContext(NamedTuple):
+  parent_graph: Any
+  device_functions: Any
+  colocation_stack: Any
+  in_cross_replica_context: Any
+  variable_policy: Any
+  xla_context_id: Any
+
+
+def make_function_context(scope_type=None) -> function_cache.FunctionContext:
+  """Generates a FunctionContext based on current contextual info."""
+  ctx = context.context()
+
+  # Don't need to open an init_scope if the tf.function call is in eager mode
+  # already.
+  executing_eagerly = ctx.executing_eagerly()
+  parent_graph = None
+  xla_context_id = 0
+  if not executing_eagerly:
+    # We want to force function retracing for each different
+    # XLAControlFlowContext, so add `xla_context_id` to the context.
+    xla_context = _enclosing_xla_context()
+    if xla_context is not None and xla_context.RequiresUniqueFunctionRetracing(
+    ):
+      xla_context_id = id(xla_context)
+
+    with ops.init_scope():
+      # The graph, or whether we're executing eagerly, should be a part of the
+      # cache key so we don't improperly capture tensors such as variables.
+      executing_eagerly = ctx.executing_eagerly()
+      parent_graph = None if executing_eagerly else ops.get_default_graph()
+
+  # pylint: disable=protected-access
+  default_graph = ops.get_default_graph()
+  # TODO(b/117617952): The current distribution strategy will affect graph
+  # building (e.g. accessing different variables from different devices) and
+  # so requires retracing for each device.
+  strategy_stack = default_graph._distribution_strategy_stack
+  uses_distribution_strategy = (
+      strategy_stack and
+      strategy_stack[-1].strategy.extended._retrace_functions_for_each_device)
+  if executing_eagerly:
+    colocation_stack = ()
+    if uses_distribution_strategy:
+      device_functions = (pydev.merge_device(ctx.device_name),)
+    else:
+      device_functions = ()
+  else:
+    colocation_stack = tuple(default_graph._colocation_stack.peek_objs())
+    if (uses_distribution_strategy or
+        func_graph_module.device_stack_has_callable(
+            default_graph._device_function_stack)):
+      # Putting the device in the cache key ensures that call-site device
+      # annotations are respected.
+      device_functions = tuple(default_graph._device_functions_outer_to_inner)
+    else:
+      device_functions = ()
+
+  in_cross_replica_context = False
+  try:
+    in_cross_replica_context = (strategy_stack[-1].replica_context is None)  # pylint: disable=protected-access
+  except (AttributeError, IndexError):
+    pass
+
+  if save_context.in_save_context():
+    variable_policy = (
+        save_context.get_save_options().experimental_variable_policy)
+  else:
+    variable_policy = None
+
+  return function_cache.FunctionContext(
+      EagerContext(
+          parent_graph,
+          device_functions,
+          colocation_stack,
+          in_cross_replica_context,
+          variable_policy,
+          xla_context_id,
+      ),
+      scope_type,
+  )
+
+
+def _enclosing_xla_context():
+  """Returns the XLAControlFlowContext, which exists inside a tpu.rewrite()."""
+  graph = ops.get_default_graph()
+  while graph is not None:
+    # pylint: disable=protected-access
+    context_ = graph._get_control_flow_context()
+    # pylint: enable=protected-access
+    while context_ is not None:
+      if isinstance(context_, control_flow_ops.XLAControlFlowContext):
+        return context_
+      context_ = context_.outer_context
+    # This may be a FuncGraph due to defuns or v2 control flow. We need to
+    # find the original graph with the XLAControlFlowContext.
+    graph = getattr(graph, "outer_graph", None)
+  return None

tensorflow/python/eager/polymorphic_function/function_spec_test.py

@@ -0,0 +1,567 @@
+# 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 function_type_utils."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.python.eager.polymorphic_function import function_type_utils
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.platform import test
+from tensorflow.python.util import tf_decorator
+
+
+def dummy_tf_decorator(func):
+
+  def wrapper(*args, **kwargs):
+    return func(*args, **kwargs)
+
+  return tf_decorator.make_decorator(func, wrapper)
+
+
+def transparent_decorator(func):
+  return func
+
+
+class FunctionSpecTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (trace_type.from_value([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_required_only(self, input_signature, type_constraint, decorator):
+
+    @decorator
+    def foo(x, y, z):  # pylint: disable=unused-argument
+      pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec),
+        (['x', 'y', 'z'], None, None, None, [], None, {}))
+    self.assertEqual(spec.input_signature, input_signature)
+    self.assertEqual(spec.default_values, {})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[0]),
+            function_type_lib.Parameter(
+                'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[1]),
+            function_type_lib.Parameter(
+                'z', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[2])
+        ]))
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint':
+              (tensor_spec.TensorSpec(shape=None),
+               tensor_spec.TensorSpec(shape=None), trace_type.from_value(3))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),),
+          'type_constraint':
+              (tensor_spec.TensorSpec(shape=None), trace_type.from_value(2),
+               trace_type.from_value(3))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (trace_type.from_value([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_optional_only(self, input_signature, type_constraint, decorator):
+
+    @decorator
+    def foo(x=1, y=2, z=3):  # pylint: disable=unused-argument
+      pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec),
+        (['x', 'y', 'z'], None, None, (1, 2, 3), [], None, {}))
+    self.assertEqual(spec.input_signature, input_signature)
+    self.assertEqual(spec.default_values, {'x': 1, 'y': 2, 'z': 3})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
+                type_constraint[0]),
+            function_type_lib.Parameter(
+                'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
+                type_constraint[1]),
+            function_type_lib.Parameter(
+                'z', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
+                type_constraint[2])
+        ]))
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint':
+              (tensor_spec.TensorSpec(shape=None),
+               tensor_spec.TensorSpec(shape=None), trace_type.from_value(3))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (trace_type.from_value([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_required_and_optional(self, input_signature, type_constraint,
+                                 decorator):
+
+    @decorator
+    def foo(x, y, z=3):  # pylint: disable=unused-argument
+      pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec),
+        (['x', 'y', 'z'], None, None, (3,), [], None, {}))
+    self.assertEqual(spec.input_signature, input_signature)
+    self.assertEqual(spec.default_values, {'z': 3})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[0]),
+            function_type_lib.Parameter(
+                'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[1]),
+            function_type_lib.Parameter(
+                'z', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
+                type_constraint[2])
+        ]))
+
+  @parameterized.product(
+      ({
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (trace_type.from_value([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_varargs(self, input_signature, type_constraint, decorator):
+
+    @decorator
+    def foo(*my_var_args):  # pylint: disable=unused-argument
+      pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec),
+        (['my_var_args_0', 'my_var_args_1', 'my_var_args_2'
+         ], None, None, None, [], None, {}))
+    self.assertEqual(spec.input_signature, input_signature)
+    self.assertEqual(spec.default_values, {})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'my_var_args_0', function_type_lib.Parameter.POSITIONAL_ONLY,
+                False, type_constraint[0]),
+            function_type_lib.Parameter(
+                'my_var_args_1', function_type_lib.Parameter.POSITIONAL_ONLY,
+                False, type_constraint[1]),
+            function_type_lib.Parameter(
+                'my_var_args_2', function_type_lib.Parameter.POSITIONAL_ONLY,
+                False, type_constraint[2])
+        ]))
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint':
+              (tensor_spec.TensorSpec(shape=None),
+               tensor_spec.TensorSpec(shape=None), trace_type.from_value(3))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (trace_type.from_value([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+                              tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_kwonly(self, input_signature, type_constraint, decorator):
+
+    @decorator
+    def foo(x, y, *, z=3):  # pylint: disable=unused-argument
+      pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec), (['x', 'y'], None, None, None, ['z'], {
+            'z': 3
+        }, {}))
+    self.assertEqual(spec.input_signature, input_signature)
+    self.assertEqual(spec.default_values, {'z': 3})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[0]),
+            function_type_lib.Parameter(
+                'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[1]),
+            function_type_lib.Parameter(
+                'z', function_type_lib.Parameter.KEYWORD_ONLY, True,
+                type_constraint[2])
+        ]))
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),),
+          'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
+                              trace_type.from_value(1))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None)),
+          'type_constraint':
+              (None,
+               trace_type.from_value([
+                   tensor_spec.TensorSpec(shape=None),
+                   tensor_spec.TensorSpec(shape=None)
+               ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+               tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_method_bound_internal(
+      self, input_signature, type_constraint, decorator
+  ):
+
+    def testing_decorator(func):
+      spec = function_type_utils.FunctionSpec.from_function_and_signature(
+          func, input_signature
+      )
+      self.assertEqual(
+          tuple(spec.fullargspec),
+          (['self', 'x', 'y'], None, None, (1,), [], None, {}),
+      )
+
+      self.assertEqual(spec.default_values, {'y': 1})
+
+      self.assertEqual(
+          spec.function_type,
+          function_type_lib.FunctionType([
+              function_type_lib.Parameter(
+                  'self',
+                  function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
+                  False,
+                  type_constraint[0],
+              ),
+              function_type_lib.Parameter(
+                  'x',
+                  function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
+                  False,
+                  type_constraint[1],
+              ),
+              function_type_lib.Parameter(
+                  'y',
+                  function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
+                  True,
+                  type_constraint[2],
+              ),
+          ]),
+      )
+
+      return func
+
+    class MyClass:
+
+      @testing_decorator
+      def foo(self, x, y=1):
+        pass
+
+    MyClass().foo(1)
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),),
+          'type_constraint': (tensor_spec.TensorSpec(shape=None),
+                              trace_type.from_value(1))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (trace_type.from_value([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+                              tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_method_bound_external(
+      self, input_signature, type_constraint, decorator
+  ):
+    class MyClass:
+
+      @decorator
+      def foo(self, x, y=1):
+        pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        MyClass().foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec),
+        (['x', 'y'], None, None, (1,), [], None, {}),
+    )
+    self.assertEqual(spec.default_values, {'y': 1})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[0]),
+            function_type_lib.Parameter(
+                'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
+                type_constraint[1])
+        ]))
+
+  @parameterized.product(
+      ({
+          'input_signature': None,
+          'type_constraint': (None, None, None)
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None)),
+          'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
+                              tensor_spec.TensorSpec(shape=None))
+      }, {
+          'input_signature': (tensor_spec.TensorSpec(shape=None),),
+          'type_constraint': (None, tensor_spec.TensorSpec(shape=None),
+                              trace_type.from_value(1))
+      }, {
+          'input_signature': ([
+              tensor_spec.TensorSpec(shape=None),
+              tensor_spec.TensorSpec(shape=None)
+          ], tensor_spec.TensorSpec(shape=None)),
+          'type_constraint':
+              (None,
+               trace_type.from_value([
+                   tensor_spec.TensorSpec(shape=None),
+                   tensor_spec.TensorSpec(shape=None)
+               ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+               tensor_spec.TensorSpec(shape=None))
+      }),
+      decorator=(dummy_tf_decorator, transparent_decorator),
+  )
+  def test_method_unbound(self, input_signature, type_constraint, decorator):
+
+    class MyClass:
+
+      @decorator
+      def foo(self, x, y=1):
+        pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        MyClass.foo, input_signature)
+    self.assertEqual(
+        tuple(spec.fullargspec),
+        (['self', 'x', 'y'], None, None, (1,), [], None, {}))
+    self.assertEqual(spec.input_signature, input_signature)
+    self.assertEqual(spec.default_values, {'y': 1})
+    self.assertEqual(
+        spec.function_type,
+        function_type_lib.FunctionType([
+            function_type_lib.Parameter(
+                'self', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
+                False, type_constraint[0]),
+            function_type_lib.Parameter(
+                'x', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, False,
+                type_constraint[1]),
+            function_type_lib.Parameter(
+                'y', function_type_lib.Parameter.POSITIONAL_OR_KEYWORD, True,
+                type_constraint[2])
+        ]))
+
+  def test_spec_summary(self):
+    input_signature = (
+        tensor_spec.TensorSpec(shape=None),
+        tensor_spec.TensorSpec(shape=None),
+    )
+
+    @dummy_tf_decorator
+    def foo(x=2, y=3):  # pylint: disable=unused-argument
+      pass
+
+    spec = function_type_utils.FunctionSpec.from_function_and_signature(
+        foo, input_signature
+    )
+    self.assertEqual(
+        spec.signature_summary(True),
+        'Input Parameters:\n'
+        + '  x (POSITIONAL_OR_KEYWORD):'
+        ' TensorSpec(shape=<unknown>, dtype=tf.float32, name=None)\n'
+        + '  y'
+        ' (POSITIONAL_OR_KEYWORD): TensorSpec(shape=<unknown>,'
+        ' dtype=tf.float32, name=None)\n'
+        + 'Output Type:\n'
+        + '  None\n'
+        + 'Captures:\n'
+        + '  None\n'
+        + 'Defaults:\n'
+        + '  x: 2\n'
+        + '  y: 3',
+    )
+
+
+# TODO(fmuham): Remove when is_same_structure is removed.
+class SameStructureTest(test.TestCase):
+
+  def test_same_structure(self):
+    self.assertTrue(
+        function_type_utils.is_same_structure([1, 2, 3], [1, 2, 3], True)
+    )
+    self.assertTrue(
+        function_type_utils.is_same_structure([1, 2, 3], [1, 2, 4], False)
+    )
+
+    self.assertFalse(
+        function_type_utils.is_same_structure([1, 2, 3], [1, 2, 4], True)
+    )
+    self.assertFalse(
+        function_type_utils.is_same_structure([1, 2, 3], [1, 2, 3, 4], False)
+    )
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/polymorphic_function/function_type_utils.py

@@ -0,0 +1,548 @@
+# 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.
+# ==============================================================================
+"""Utilities for using FunctionType with tf.function."""
+
+import functools
+import inspect
+from typing import Any, Dict, Tuple
+
+import six
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.util import nest
+
+
+def to_fullargspec(
+    function_type: function_type_lib.FunctionType,
+    default_values: Dict[str, Any],
+) -> inspect.FullArgSpec:
+  """Generates backwards compatible FullArgSpec from FunctionType."""
+  args = []
+  varargs = None
+  varkw = None
+  defaults = []
+  kwonlyargs = []
+  kwonlydefaults = {}
+
+  for parameter in function_type.parameters.values():
+    if parameter.kind in [
+        inspect.Parameter.POSITIONAL_ONLY,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    ]:
+      args.append(parameter.name)
+      if parameter.default is not inspect.Parameter.empty:
+        defaults.append(default_values[parameter.name])
+    elif parameter.kind is inspect.Parameter.KEYWORD_ONLY:
+      kwonlyargs.append(parameter.name)
+      if parameter.default is not inspect.Parameter.empty:
+        kwonlydefaults[parameter.name] = default_values[parameter.name]
+    elif parameter.kind is inspect.Parameter.VAR_POSITIONAL:
+      varargs = parameter.name
+    elif parameter.kind is inspect.Parameter.VAR_KEYWORD:
+      varkw = parameter.name
+
+  return inspect.FullArgSpec(
+      args,
+      varargs,
+      varkw,
+      tuple(defaults) if defaults else None,
+      kwonlyargs,
+      kwonlydefaults if kwonlydefaults else None,
+      annotations={},
+  )
+
+
+def _to_default_values(fullargspec):
+  """Returns default values from the function's inspected fullargspec."""
+  if fullargspec.defaults is not None:
+    defaults = {
+        name: value
+        for name, value in zip(
+            fullargspec.args[-len(fullargspec.defaults) :], fullargspec.defaults
+        )
+    }
+  else:
+    defaults = {}
+
+  if fullargspec.kwonlydefaults is not None:
+    defaults.update(fullargspec.kwonlydefaults)
+
+  defaults = {
+      function_type_lib.sanitize_arg_name(name): value
+      for name, value in defaults.items()
+  }
+
+  return defaults
+
+
+def to_function_type(fullargspec):
+  """Generates FunctionType and default values from fullargspec."""
+  default_values = _to_default_values(fullargspec)
+  parameters = []
+
+  for arg in fullargspec.args:
+    arg_name = function_type_lib.sanitize_arg_name(arg)
+    parameters.append(
+        function_type_lib.Parameter(
+            arg_name,
+            function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
+            arg_name in default_values,
+            None,
+        )
+    )
+
+  if fullargspec.varargs is not None:
+    parameters.append(
+        function_type_lib.Parameter(
+            fullargspec.varargs,
+            function_type_lib.Parameter.VAR_POSITIONAL,
+            False,
+            None,
+        )
+    )
+
+  for kwarg in fullargspec.kwonlyargs:
+    parameters.append(
+        function_type_lib.Parameter(
+            function_type_lib.sanitize_arg_name(kwarg),
+            function_type_lib.Parameter.KEYWORD_ONLY,
+            kwarg in default_values,
+            None,
+        )
+    )
+
+  if fullargspec.varkw is not None:
+    parameters.append(
+        function_type_lib.Parameter(
+            fullargspec.varkw,
+            function_type_lib.Parameter.VAR_KEYWORD,
+            False,
+            None,
+        )
+    )
+
+  return function_type_lib.FunctionType(parameters), default_values
+
+
+def to_input_signature(function_type):
+  """Extracts an input_signature from function_type instance."""
+  constrained_parameters = list(function_type.parameters.keys())
+
+  # self does not have a constraint in input_signature
+  if "self" in constrained_parameters:
+    constrained_parameters.pop(0)
+
+  # There are no parameters to constrain.
+  if not constrained_parameters:
+    return tuple()
+
+  constraints = []
+  is_auto_constrained = False
+
+  for parameter_name in constrained_parameters:
+    parameter = function_type.parameters[parameter_name]
+    constraint = None
+    if parameter.type_constraint:
+      # Generate legacy constraint representation.
+      constraint = parameter.type_constraint.placeholder_value(
+          trace_type.InternalPlaceholderContext(unnest_only=True)
+      )
+      if any(
+          not isinstance(arg, tensor.TensorSpec)
+          for arg in nest.flatten([constraint], expand_composites=True)
+      ):
+        # input_signature only supports contiguous TensorSpec composites
+        is_auto_constrained = True
+        break
+      else:
+        constraints.append(constraint)
+
+  # All constraints were generated by FunctionType
+  if is_auto_constrained and not constraints:
+    return tuple()
+
+  # If the list is empty then there was no input_signature specified.
+  return tuple(constraints) if constraints else None
+
+
+def to_arg_names(function_type):
+  """Generates a list of arg names from a FunctionType."""
+  arg_names = []
+  for p in function_type.parameters.values():
+    if p.kind in {
+        function_type_lib.Parameter.POSITIONAL_ONLY,
+        function_type_lib.Parameter.POSITIONAL_OR_KEYWORD,
+    }:
+      arg_names.append(p.name)
+  return arg_names
+
+
+# TODO(b/214462107): Minimize API surface for FunctionSpec.
+class FunctionSpec(object):
+  """Specification of how to bind arguments to a function.
+
+  Deprecated. Please use FunctionType instead.
+  """
+
+  @classmethod
+  def from_function_and_signature(
+      cls, python_function, input_signature, is_pure=False, jit_compile=None
+  ):
+    """Creates a FunctionSpec instance given a python function and signature.
+
+    Args:
+      python_function: a function to inspect
+      input_signature: a signature of the function (None, if variable)
+      is_pure: if True all input arguments (including variables and constants)
+        will be converted to tensors and no variable changes allowed.
+      jit_compile: see `tf.function`
+
+    Returns:
+      instance of FunctionSpec
+    """
+    function_type, default_values = make_function_type(
+        python_function, input_signature)
+    # Get the function's name.  Remove functools.partial wrappers if necessary.
+    while isinstance(python_function, functools.partial):
+      python_function = python_function.func
+    name = getattr(python_function, "__name__", "f")
+
+    return FunctionSpec(
+        function_type,
+        default_values,
+        is_pure=is_pure,
+        jit_compile=jit_compile,
+        name=name,
+    )
+
+  @classmethod
+  def from_fullargspec_and_signature(
+      cls,
+      fullargspec,
+      input_signature,
+      is_pure=False,
+      name=None,
+      jit_compile=None,
+  ):
+    """Construct FunctionSpec from legacy FullArgSpec format."""
+    function_type, default_values = to_function_type(fullargspec)
+    if input_signature:
+      input_signature = tuple(input_signature)
+      _validate_signature(input_signature)
+      function_type = function_type_lib.add_type_constraints(
+          function_type, input_signature, default_values
+      )
+
+    return FunctionSpec(
+        function_type, default_values, is_pure, name, jit_compile
+    )
+
+  def __init__(
+      self,
+      function_type,
+      default_values,
+      is_pure=False,
+      name=None,
+      jit_compile=None,
+  ):
+    """Constructs a FunctionSpec describing a python function.
+
+    Args:
+      function_type: A FunctionType describing the python function signature.
+      default_values: Dictionary mapping parameter names to default values.
+      is_pure: if True all input arguments (including variables and constants)
+        will be converted to tensors and no variable changes allowed.
+      name: Name of the function
+      jit_compile: see `tf.function`.
+    """
+    self._function_type = function_type
+    self._default_values = default_values
+    self._fullargspec = to_fullargspec(function_type, default_values)
+    self._is_pure = is_pure
+    self._jit_compile = jit_compile
+
+    # TODO(edloper): Include name when serializing for SavedModel?
+    self._name = name or "f"
+    self._input_signature = to_input_signature(function_type)
+
+  @property
+  def default_values(self):
+    """Returns dict mapping parameter names to default values."""
+    return self._default_values
+
+  @property
+  def function_type(self):
+    """Returns a FunctionType representing the Python function signature."""
+    return self._function_type
+
+  @property
+  def fullargspec(self):
+    return self._fullargspec
+
+  # TODO(fmuham): Replace usages with FunctionType and remove.
+  @property
+  def input_signature(self):
+    return self._input_signature
+
+  # TODO(fmuham): Replace usages with FunctionType and remove.
+  @property
+  def flat_input_signature(self):
+    return tuple(nest.flatten(self.input_signature, expand_composites=True))
+
+  @property
+  def is_pure(self):
+    return self._is_pure
+
+  @property
+  def jit_compile(self):
+    return self._jit_compile
+
+  # TODO(fmuham): Replace usages and remove.
+  @property
+  def arg_names(self):
+    return to_arg_names(self.function_type)
+
+  def signature_summary(self, default_values=False):
+    """Returns a string summarizing this function's signature.
+
+    Args:
+      default_values: If true, then include default values in the signature.
+
+    Returns:
+      A `string`.
+    """
+    summary = f"{self._function_type!r}"
+    if default_values:
+      summary += "\nDefaults:"
+      if self.default_values:
+        for name, value in self.default_values.items():
+          summary += f"\n  {name}: {value!r}"
+      else:
+        summary += "\n  None"
+    return summary
+
+
+def make_function_type(python_function, input_signature):
+  """Generates a FunctionType for python_function."""
+  _validate_signature(input_signature)
+
+  function_type = function_type_lib.FunctionType.from_callable(
+      python_function
+  )
+  default_values = function_type_lib.FunctionType.get_default_values(
+      python_function
+  )
+
+  if input_signature is not None:
+    input_signature = tuple(input_signature)
+    function_type = function_type_lib.add_type_constraints(
+        function_type, input_signature, default_values
+    )
+
+  return function_type, default_values
+
+
+def make_canonicalized_monomorphic_type(
+    args: Any,
+    kwargs: Any,
+    capture_types: Any,
+    polymorphic_type,
+) -> Tuple[function_type_lib.FunctionType, trace_type.InternalTracingContext]:
+  """Generates function type given the function arguments."""
+  kwargs = {
+      function_type_lib.sanitize_arg_name(name): value
+      for name, value in kwargs.items()
+  }
+
+  function_type, type_context = (
+      function_type_lib.canonicalize_to_monomorphic(
+          args, kwargs, {}, capture_types, polymorphic_type
+      )
+  )
+
+  return function_type, type_context
+
+
+def canonicalize_function_inputs(
+    args, kwargs, function_type, default_values=None, is_pure=False
+):
+  """Canonicalizes `args` and `kwargs`.
+
+  Canonicalize the inputs to the Python function using FunctionType.
+  In particular, we parse the varargs and kwargs that the
+  original function was called with into a tuple corresponding to the
+  Python function's positional (named) arguments and a dictionary
+  corresponding to its kwargs.  Missing default arguments are added.
+
+  If the FunctionType has an type constraints, then they are used to convert
+  arguments to tensors; otherwise, any inputs containing numpy arrays are
+  converted to tensors.
+
+
+  Args:
+    args: The varargs this object was called with.
+    kwargs: The keyword args this function was called with.
+    function_type: FunctionType to canonicalize against.
+    default_values: Default values to use.
+    is_pure: Force variable inputs to Tensors.
+
+  Returns:
+    A canonicalized ordering of the inputs, as well as full and filtered
+    (Tensors and Variables only) versions of their concatenated flattened
+    representations, represented by a tuple in the form (args, kwargs,
+    flat_args, filtered_flat_args). Here: `args` is a full list of bound
+    arguments, and `kwargs` contains only true keyword arguments, as opposed
+    to named arguments called in a keyword-like fashion.
+
+  Raises:
+    ValueError: If a keyword in `kwargs` cannot be matched with a positional
+      argument when an input signature is specified, or when the inputs
+      do not conform to the input signature.
+  """
+  default_values = {} if not default_values else default_values
+  if is_pure:
+    args, kwargs = _convert_variables_to_tensors(args, kwargs)
+  bound_arguments = bind_function_inputs(
+      args, kwargs, function_type, default_values
+  )
+  return bound_arguments
+
+
+def bind_function_inputs(args, kwargs, function_type, default_values):
+  """Bind `args` and `kwargs` into a canonicalized signature args, kwargs."""
+  sanitized_kwargs = {
+      function_type_lib.sanitize_arg_name(k): v for k, v in kwargs.items()
+  }
+  if len(kwargs) != len(sanitized_kwargs):
+    raise ValueError(
+        "Name collision after sanitization. Please rename "
+        "tf.function input parameters. Original: "
+        f"{sorted(kwargs.keys())}, Sanitized: "
+        f"{sorted(sanitized_kwargs.keys())}"
+    )
+
+  try:
+    bound_arguments = function_type.bind_with_defaults(
+        args, sanitized_kwargs, default_values
+    )
+  except Exception as e:
+    raise TypeError(
+        f"Binding inputs to tf.function failed due to `{e}`. "
+        f"Received args: {args} and kwargs: {sanitized_kwargs} for signature:"
+        f" {function_type}."
+    ) from e
+  return bound_arguments
+
+
+def _validate_signature(signature):
+  """Checks the input_signature to be valid."""
+  if signature is None:
+    return
+
+  if not isinstance(signature, (tuple, list)):
+    raise TypeError(
+        "input_signature must be either a tuple or a list, got "
+        f"{type(signature)}."
+    )
+
+  # TODO(xjun): Allow VariableSpec once we figure out API for de-aliasing.
+  variable_specs = _get_variable_specs(signature)
+  if variable_specs:
+    raise TypeError(
+        f"input_signature doesn't support VariableSpec, got {variable_specs}"
+    )
+
+  if any(
+      not isinstance(arg, tensor.TensorSpec)
+      for arg in nest.flatten(signature, expand_composites=True)
+  ):
+    bad_args = [
+        arg
+        for arg in nest.flatten(signature, expand_composites=True)
+        if not isinstance(arg, tensor.TensorSpec)
+    ]
+    raise TypeError(
+        "input_signature must be a possibly nested sequence of "
+        f"TensorSpec objects, got invalid args {bad_args} with "
+        f"types {list(six.moves.map(type, bad_args))}."
+    )
+
+
+def _to_tensor_or_tensor_spec(x):
+  return (
+      x
+      if isinstance(x, (tensor.Tensor, tensor.TensorSpec))
+      else ops.convert_to_tensor(x)
+  )
+
+
+def _convert_variables_to_tensors(args, kwargs):
+  args = [_to_tensor_or_tensor_spec(x) for x in args]
+  kwargs = {kw: _to_tensor_or_tensor_spec(x) for kw, x in kwargs.items()}
+  return tuple(args), kwargs
+
+
+def _get_variable_specs(args):
+  """Returns `VariableSpecs` from `args`."""
+  variable_specs = []
+  for arg in nest.flatten(args):
+    if not isinstance(arg, type_spec.TypeSpec):
+      continue
+    if isinstance(arg, resource_variable_ops.VariableSpec):
+      variable_specs.append(arg)
+    elif not isinstance(arg, tensor.TensorSpec):
+      # arg is a CompositeTensor spec.
+      variable_specs.extend(_get_variable_specs(arg._component_specs))  # pylint: disable=protected-access
+  return variable_specs
+
+
+def derive_from_graph(func_graph):
+  """Derives a FunctionType from FuncGraph."""
+  # TODO(fmuham): Include structure info from structured_inputs
+  input_signature = (
+      tuple(trace_type.from_value(i) for i in func_graph.inputs),
+      {},
+  )
+
+  # TODO(fmuham): Include output structure info from structured_outputs
+  output_signature = tuple(trace_type.from_value(o) for o in func_graph.outputs)
+
+  return function_type_lib.from_structured_signature(
+      input_signature,
+      output_signature,
+      func_graph.function_captures.capture_types,
+  )
+
+
+# TODO(fmuham): Replace usages with TraceType and remove.
+def is_same_structure(structure1, structure2, check_values=False):
+  """Check two structures for equality, optionally of types and of values."""
+  try:
+    nest.assert_same_structure(structure1, structure2, expand_composites=True)
+  except (ValueError, TypeError):
+    return False
+  if check_values:
+    flattened1 = nest.flatten(structure1, expand_composites=True)
+    flattened2 = nest.flatten(structure2, expand_composites=True)
+    # First check the types to avoid AttributeErrors.
+    if any(type(f1) is not type(f2) for f1, f2 in zip(flattened1, flattened2)):
+      return False
+    return flattened1 == flattened2
+  return True

tensorflow/python/eager/polymorphic_function/gradients_test.py

@@ -0,0 +1,927 @@
+# Copyright 2017 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_grad
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+
+
+_COS_DERIVATIVES = [math_ops.cos,
+                    lambda x: -math_ops.sin(x),
+                    lambda x: -math_ops.cos(x),
+                    math_ops.sin,
+                    math_ops.cos]
+
+
+class FunctionGradientsTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(FunctionGradientsTest, self).setUp()
+    cpus = config.list_physical_devices('CPU')
+    # Set 4 virtual CPUs
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+
+  def testGraphModeWithGradients(self):
+    v = resource_variable_ops.ResourceVariable(1.0, name='v')
+
+    @polymorphic_function.function
+    def step():
+      def inner():
+        return v * v
+
+      return backprop.implicit_grad(inner)()[0][0]
+
+    self.assertAllEqual(step(), 2.0)
+
+  def testGraphGradientVariable(self):
+    with ops.Graph().as_default(), self.cached_session():
+      v = variables.Variable(1.0)
+
+      @polymorphic_function.function
+      def f():
+        return 2.0 * v
+
+      node = f()
+      grads, = gradients_impl.gradients(node, v)
+      v.initializer.run()
+      self.assertAllEqual(grads, 2.0)
+      self.assertEqual(grads.shape, v.shape)
+
+  def testSymbolicHigherOrder(self):
+    @polymorphic_function.function
+    def f(x, order):
+      y = polymorphic_function.function(lambda: math_ops.cos(x))()
+      for _ in range(order):
+        y, = gradients_impl.gradients(y, [x])
+      return y
+    for order, expected in enumerate(_COS_DERIVATIVES):
+      self.assertAllClose(
+          expected(constant_op.constant(1.)),
+          f(constant_op.constant(1.), order))
+
+  @parameterized.parameters([dict(persistent=True),
+                             dict(persistent=False)])
+  def testSymbolicHigherOrderUnderTape(self, persistent):
+    @polymorphic_function.function
+    def f(x, order):
+      with backprop.GradientTape(persistent=persistent) as tape:
+        tape.watch(x)
+        # Note that having a tape active, even if we don't use it, forces us
+        # down a different function call path. Symbolic gradients should work
+        # here too; correctness of tape gradients are tested elsewhere.
+        y = polymorphic_function.function(lambda: math_ops.cos(x))()
+      tape_dy = tape.gradient(y, x)
+      for _ in range(order):
+        y, = gradients_impl.gradients(y, [x])
+      if order > 0:
+        y1 = tape_dy
+        for _ in range(order - 1):
+          y1, = gradients_impl.gradients(y1, [x])
+      else:
+        y1 = y
+      return y, y1
+    for order, expected_f in enumerate(_COS_DERIVATIVES):
+      expected = self.evaluate(expected_f(constant_op.constant(1.)))
+      self.assertAllClose(
+          (expected, expected),
+          f(constant_op.constant(1.), order))
+
+  def testIteratedGradientsNested(self):
+
+    def _grad(f):
+      def _grad_function(primal):
+        with backprop.GradientTape() as tape:
+          tape.watch(primal)
+          primal_out = f(primal)
+        return tape.gradient(primal_out, primal)
+      return _grad_function
+
+    @polymorphic_function.function
+    def _forward(x):
+      return math_ops.cos(x)
+
+    f = _forward
+    traced_f = polymorphic_function.function(f)
+    one = constant_op.constant(1.)
+    for expected in _COS_DERIVATIVES:
+      self.assertAllClose(expected(one), f(one))
+      self.assertAllClose(expected(one), traced_f(one))
+      self.assertAllClose(expected(one), polymorphic_function.function(f)(one))
+      f = _grad(f)
+      traced_f = polymorphic_function.function(_grad(traced_f))
+
+  def testIteratedGradientsNestedWithVariable(self):
+
+    def _grad(f):
+      def _grad_function():
+        with backprop.GradientTape() as tape:
+          primal_out = f()
+        g, = tape.gradient(primal_out, tape.watched_variables())
+        return g
+      return _grad_function
+
+    v = variables.Variable(2.)
+
+    @polymorphic_function.function
+    def _forward():
+      return math_ops.cos(v)
+
+    f = _forward
+
+    two = constant_op.constant(2.)
+
+    for expected in _COS_DERIVATIVES:
+      self.assertAllClose(expected(two), f())
+      self.assertAllClose(expected(two), polymorphic_function.function(f)())
+      f = _grad(f)
+
+  def testIteratedGradientsPersistent(self):
+
+    @polymorphic_function.function
+    def _forward(z):
+      return math_ops.cos(z)
+
+    f = _forward
+    with backprop.GradientTape(persistent=True) as tape:
+      start = constant_op.constant(1.)
+      tape.watch(start)
+      x = f(start)
+      for expected in _COS_DERIVATIVES:
+        self.assertAllClose(expected(start), x)
+        x = tape.gradient(x, start)
+
+  def testHigherOrderWithVariable(self):
+
+    v = variables.Variable(1.)
+
+    @polymorphic_function.function
+    def _forward():
+      return math_ops.cos(v)
+
+    f = _forward
+    with backprop.GradientTape(persistent=True) as tape:
+      x = f()
+      for expected in _COS_DERIVATIVES:
+        self.assertAllClose(expected(constant_op.constant(1.)), x)
+        x, = tape.gradient(x, tape.watched_variables())
+
+  def testGradientsChained(self):
+
+    @polymorphic_function.function
+    def _forward(z):
+      return math_ops.cos(z)
+
+    f = _forward
+    x = constant_op.constant(1.)
+    with backprop.GradientTape() as t:
+      t.watch(x)
+      y = f(x)
+    with backprop.GradientTape() as tt:
+      doutputs = constant_op.constant(2.)
+      tt.watch(doutputs)
+      g = t.gradient(y, x, doutputs)
+    self.assertAllClose(-2. * math_ops.sin(x), g)
+    gg = tt.gradient(g, doutputs)
+    # We're taking gradients with respect to doutputs, which is just a linear
+    # function of the gradient.
+    self.assertAllClose(-math_ops.sin(x), gg)
+
+  def testSymGradGatherNd(self):
+    with ops.Graph().as_default(), self.cached_session():
+
+      @polymorphic_function.function
+      def f(x):
+        return array_ops.gather_nd(x, [[0]])
+
+      c = constant_op.constant([[2.]])
+      f_c = f(c)
+      g, = gradients_impl.gradients(f_c, c)
+      self.assertAllEqual(self.evaluate(g).values, [[1.0]])
+
+  def testNoSymGradNestedDefun(self):
+
+    @polymorphic_function.function
+    def outer():
+
+      @polymorphic_function.function
+      def f(x):
+        return array_ops.gather_nd(x, [[0]])
+
+      c = constant_op.constant([[2.]])
+      f_c = f(c)
+      g, = gradients_impl.gradients(f_c, c)
+      self.assertIsInstance(g, indexed_slices.IndexedSlices)
+
+    outer()
+
+  def testGraphFunctionWithGradients(self):
+    v = resource_variable_ops.ResourceVariable(1.0, name='v')
+
+    @polymorphic_function.function
+    def step():
+      def inner():
+        return v * v
+
+      return backprop.implicit_grad(inner)()[0][0]
+
+    step_op = step.get_concrete_function()
+    self.assertEqual(step_op.output_dtypes, dtypes.float32)
+    self.assertEqual(step_op.output_shapes, tensor_shape.TensorShape([]))
+    self.assertAllEqual(step_op(), 2.0)
+
+  @test_util.run_in_graph_and_eager_modes()
+  def testDefunCondGradient(self):
+
+    @polymorphic_function.function
+    def f(x):
+      return cond.cond(x > 0.5, lambda: 2 * x, lambda: 3 * x)
+
+    with backprop.GradientTape() as t:
+      x = constant_op.constant(1.0)
+      t.watch(x)
+      y = f(x)
+    self.assertAllEqual(self.evaluate(t.gradient(y, x)), 2.0)
+
+  @test_util.run_in_graph_and_eager_modes()
+  def testGraphLoopGradient(self):
+
+    @polymorphic_function.function
+    def f(x):
+      return while_loop.while_loop(
+          lambda _, i: i < 2, lambda x, i: (2 * x, i + 1), [x, 0]
+      )[0]
+
+    with backprop.GradientTape() as t:
+      x = constant_op.constant(1.0)
+      t.watch(x)
+      y = f(x)
+    self.assertAllEqual(self.evaluate(t.gradient(y, x)), 4.0)
+
+  def testGraphLoopGradientInsideSession(self):
+    with ops.Graph().as_default():
+      n = constant_op.constant(2.0)
+      x = array_ops.placeholder(dtypes.float32, shape=None)
+
+      @polymorphic_function.function
+      def f():
+        c = lambda n: n < 10
+        b = lambda n: n * x
+        return while_loop.while_loop(c, b, [n], [tensor_shape.unknown_shape()])
+
+      l = f()
+      dx = gradients_impl.gradients(l, [x])[0]
+
+      with self.cached_session():
+        self.assertEqual(dx.eval(feed_dict={x: 2.0}), 24.0)
+
+  def testDefunDifferentiable(self):
+    v = resource_variable_ops.ResourceVariable(1.0)
+
+    @polymorphic_function.function
+    def f():
+      return v * v
+
+    self.assertAllEqual(backprop.implicit_grad(f)()[0][0], 2.0)
+
+  def testDefunCanBeDifferentiatedTwice(self):
+    v = resource_variable_ops.ResourceVariable(1.0)
+
+    @polymorphic_function.function
+    def f():
+      return v * v
+
+    self.assertAllEqual(backprop.implicit_grad(f)()[0][0], 2.0)
+    # Ensure that v is watched again.
+    self.assertAllEqual(backprop.implicit_grad(f)()[0][0], 2.0)
+
+  def testSymbolicGradientVariableNoneNotZerosLike(self):
+    with ops.Graph().as_default():
+      v = variables.Variable(1.0)
+
+      @polymorphic_function.function
+      def f(x, v):
+        v.read_value()
+        return x * x
+
+      x = constant_op.constant(1.0)
+      l = f(x, v)
+      _, dv = gradients_impl.gradients(l, [x, v])
+      with self.cached_session():
+        v.initializer.run()
+        self.assertEqual(dv, None)
+
+  def testDefunCallBackprop(self):
+
+    @polymorphic_function.function
+    def f(x):
+      return math_ops.add(x, x)
+
+    @polymorphic_function.function
+    def g(x):
+      return backprop.gradients_function(f, [0])(x)[0]
+
+    self.assertAllEqual(2, g(constant_op.constant(2.)))
+
+  @test_util.run_v1_only('b/120545219')
+  def testGraphModeEagerGradError(self):
+    with context.graph_mode():
+      def f():
+        x = variable_scope.get_variable(
+            'v', initializer=constant_op.constant(1.0))
+        return x * constant_op.constant(2.0)
+
+      with self.assertRaisesRegex(ValueError,
+                                  'No trainable variables were accessed'):
+        backprop.implicit_val_and_grad(f)()
+
+  def testDefunCallBackpropUsingSameObjectForMultipleArguments(self):
+
+    @polymorphic_function.function
+    def g(x):
+      return backprop.gradients_function(math_ops.multiply, [0, 1])(x, x)
+
+    def np_g(x):
+      return [d.numpy() for d in g(x)]
+
+    x = constant_op.constant(1.)
+    self.assertAllEqual([1., 1.], np_g(x))
+    self.assertAllEqual([1., 1.], np_g(1.))
+
+  def testGradientTensorConversionWithDefun(self):
+    three = resource_variable_ops.ResourceVariable(3.0, name='v')
+
+    @polymorphic_function.function
+    def f(x):
+      return math_ops.add(x, three)
+
+    def g(x):
+      return f(x)
+
+    g = backprop.implicit_grad(g)(constant_op.constant(1.0))[0][0]
+    self.assertAllEqual(g, 1.0)
+
+  def testGradient(self):
+    matmul = polymorphic_function.function(math_ops.matmul)
+
+    def sq(x):
+      return matmul(x, x, transpose_a=True)
+
+    t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+    grad_t, = backprop.gradients_function(sq, [0])(t)
+    self.assertAllEqual(grad_t, [[6, 6], [14, 14]])
+
+  def testGradientInFunction(self):
+
+    @polymorphic_function.function
+    def f(x):
+      return backprop.gradients_function(lambda y: y * y, [0])(x)[0]
+
+    self.assertAllEqual(f(constant_op.constant(1.0)), 2.0)
+
+  def testGradientOfGatherWithDefun(self):
+    v = resource_variable_ops.ResourceVariable([0.0, 1.0, 2.0])
+
+    def sum_gather():
+      return math_ops.reduce_sum(array_ops.gather(v, [1, 2]))
+
+    grad_fn = backprop.implicit_grad(sum_gather)
+    gradient = grad_fn()
+    defun_grad_fn = backprop.implicit_grad(
+        polymorphic_function.function(sum_gather))
+    defun_gradient = defun_grad_fn()
+    self.assertEqual(len(gradient), len(defun_gradient))
+
+    gradient = gradient[0][0]
+    defun_gradient = defun_gradient[0][0]
+    self.assertAllEqual(gradient.values, defun_gradient.values)
+    self.assertAllEqual(gradient.indices, defun_gradient.indices)
+    self.assertAllEqual(gradient.dense_shape, defun_gradient.dense_shape)
+
+  def testDifferentiableFunctionNoneOutputs(self):
+
+    @polymorphic_function.function
+    def my_function(x):
+      return x, None
+
+    def wrapper(x):
+      return my_function(x)[0]
+
+    g = backprop.gradients_function(wrapper, [0])(constant_op.constant(0.0))
+    self.assertAllEqual(g[0], 1.)
+
+    @polymorphic_function.function
+    def foo(a):
+      return None, a * a
+
+    x = constant_op.constant(5.0)
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      none, r = foo(x)
+    g = tp.gradient(r, x)
+
+    self.assertIs(none, None)
+    self.assertAllEqual(r, 25.0)
+    self.assertAllEqual(g, 2 * 5.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNestedDifferentiableFunction(self):
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return a * math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return inner_fn(x, 1.0)
+
+    x = constant_op.constant(5.0)
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunction(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return middle_fn(x, 1.0)
+
+    x = constant_op.constant(5.0)
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionWithMultipleGradCalls(self):
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return math_ops.mul(a, inner_fn(a, b))
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return middle_fn(x, 3.0)
+
+    x = constant_op.constant(5.0)
+    self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
+
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 3.0)
+    self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
+    self.assertAllEqual(middle_fn(3.0, x), 3.0 * (3.0 + 5.0))
+
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 3.0)
+
+    y = constant_op.constant(4.0)
+    with backprop.GradientTape() as tp:
+      tp.watch(y)
+      result = outer_fn(y)
+    grad = tp.gradient(result, y)
+
+    self.assertAllEqual(grad, 2 * 4.0 + 3.0)
+
+    with backprop.GradientTape() as tp:
+      tp.watch(y)
+      result = inner_fn(y, y)
+    grad = tp.gradient(result, y)
+
+    self.assertAllEqual(grad, 2.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionGradientTapeInDefun(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      with backprop.GradientTape() as tp:
+        tp.watch(x)
+        result = middle_fn(x, 1.0)
+      grad = tp.gradient(result, x)
+      return grad
+
+    x = constant_op.constant(5.0)
+    grad = outer_fn(x)
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionGradientTapeInNestedDefun(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def almost_outer_fn(x):
+      with backprop.GradientTape() as tp:
+        tp.watch(x)
+        result = middle_fn(x, 1.0)
+      grad = tp.gradient(result, x)
+      return grad
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return almost_outer_fn(x)
+
+    x = constant_op.constant(5.0)
+    grad = outer_fn(x)
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionGradientTapeInMultNestedDefun(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def almost_outer_fn(x):
+      with backprop.GradientTape() as tp:
+        tp.watch(x)
+        result = middle_fn(x, 1.0)
+      grad = tp.gradient(result, x)
+      return grad
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return almost_outer_fn(x)
+
+    @polymorphic_function.function
+    def outer_outer_fn(x):
+      return outer_fn(x)
+
+    x = constant_op.constant(5.0)
+    grad = outer_outer_fn(x)
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionTFGradientInDefun(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      result = middle_fn(x, 1.0)
+      return gradients_impl.gradients(result, [x])[0]
+
+    x = constant_op.constant(5.0)
+    grad = outer_fn(x)
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionTFGradientInNestedDefun(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def almost_outer_fn(x):
+      result = middle_fn(x, 1.0)
+      return gradients_impl.gradients(result, [x])[0]
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return almost_outer_fn(x)
+
+    x = constant_op.constant(5.0)
+    grad = outer_fn(x)
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDeeplyNestedDifferentiableFunctionTFGradientInMultNestedDefun(self):
+    @polymorphic_function.function
+    def inner_inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return inner_inner_fn(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def almost_outer_fn(x):
+      result = middle_fn(x, 1.0)
+      return gradients_impl.gradients(result, [x])[0]
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return almost_outer_fn(x)
+
+    @polymorphic_function.function
+    def outer_outer_fn(x):
+      return outer_fn(x)
+
+    x = constant_op.constant(5.0)
+    grad = outer_outer_fn(x)
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  def testDeeplyNestedDifferentiableFunctionWithVariable(self):
+    var = variables.Variable(constant_op.constant(1.0))
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return a * inner_fn(a, b)
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return middle_fn(x, var)
+
+    x = constant_op.constant(5.0)
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 1.0)
+
+  def testDeeplyNestedDifferentiableFunctionWithVariableMultipleGradCalls(self):
+    v = variables.Variable(constant_op.constant(3.0))
+
+    @polymorphic_function.function
+    def inner_fn(a, b):
+      return math_ops.add(a, b)
+
+    @polymorphic_function.function
+    def middle_fn(a, b):
+      return math_ops.mul(a, inner_fn(a, b))
+
+    @polymorphic_function.function
+    def outer_fn(x):
+      return middle_fn(x, v)
+
+    x = constant_op.constant(5.0)
+    self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
+
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 3.0)
+    self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
+    self.assertAllEqual(middle_fn(v, x), 3.0 * (3.0 + 5.0))
+
+    with backprop.GradientTape() as tp:
+      tp.watch(x)
+      result = outer_fn(x)
+    grad = tp.gradient(result, x)
+
+    self.assertAllEqual(grad, 2 * 5.0 + 3.0)
+
+    y = constant_op.constant(4.0)
+    with backprop.GradientTape() as tp:
+      tp.watch(y)
+      result = outer_fn(y)
+    grad = tp.gradient(result, y)
+
+    self.assertAllEqual(grad, 2 * 4.0 + 3.0)
+
+    v.assign(constant_op.constant(1.5))
+    with backprop.GradientTape() as tp:
+      tp.watch(y)
+      result = outer_fn(y)
+    grad = tp.gradient(result, y)
+
+    self.assertAllEqual(grad, 2 * 4.0 + 1.5)
+
+    with backprop.GradientTape() as tp:
+      tp.watch(y)
+      result = inner_fn(y, v)
+    grad = tp.gradient(result, y)
+
+    self.assertAllEqual(grad, 1.0)
+
+  def testDeeplyNestedDifferentiableFunctionWithVariableMultipleTFGrads(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(3.0)
+      v.initializer.run()
+
+      @polymorphic_function.function
+      def inner_fn(a, b):
+        return math_ops.add(a, b)
+
+      @polymorphic_function.function
+      def middle_fn(a, b):
+        return math_ops.mul(a, inner_fn(a, b))
+
+      @polymorphic_function.function
+      def outer_fn(x):
+        return middle_fn(x, v)
+
+      x = constant_op.constant(5.0)
+      self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
+
+      grad, = gradients_impl.gradients(outer_fn(x), x)
+
+      self.assertAllEqual(grad, 2 * 5.0 + 3.0)
+      self.assertAllEqual(outer_fn(x), 5.0 * (5.0 + 3.0))
+      self.assertAllEqual(middle_fn(v, x), 3.0 * (3.0 + 5.0))
+
+      grad, = gradients_impl.gradients(outer_fn(x), x)
+
+      self.assertAllEqual(grad, 2 * 5.0 + 3.0)
+
+      y = constant_op.constant(4.0)
+      grad, = gradients_impl.gradients(outer_fn(y), y)
+      self.assertAllEqual(grad, 2 * 4.0 + 3.0)
+
+      self.evaluate(v.assign(constant_op.constant(1.5)))
+      grad, = gradients_impl.gradients(outer_fn(y), y)
+
+      self.assertAllEqual(grad, 2 * 4.0 + 1.5)
+
+      grad, = gradients_impl.gradients(inner_fn(y, v), y)
+      self.assertAllEqual(grad, 1.0)
+
+  def testNestedDifferentiableFunctionNoneOutputs(self):
+    @polymorphic_function.function
+    def foo(a, b):
+      return None, a * math_ops.add(a, b), None, 2*a
+
+    @polymorphic_function.function
+    def bar(x):
+      return foo(x, 1.0)
+
+    x = constant_op.constant(5.0)
+    with backprop.GradientTape(persistent=True) as tp:
+      tp.watch(x)
+      none1, r1, none2, r2 = bar(x)
+    g1 = tp.gradient(r1, x)
+    g2 = tp.gradient(r2, x)
+
+    self.assertAllEqual(r1, 30.0)
+    self.assertAllEqual(r2, 10.0)
+    self.assertIs(none1, None)
+    self.assertIs(none2, None)
+    self.assertAllEqual(g1, 2 * 5.0 + 1.0)
+    self.assertAllEqual(g2, 2.0)
+
+  def testGradientWithKeywordArguments(self):
+    matmul = polymorphic_function.function(math_ops.matmul)
+
+    def sq(x):
+      return matmul(a=x, b=x, transpose_a=True)
+
+    t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+    grad_t, = backprop.gradients_function(sq, [0])(t)
+    self.assertAllEqual(grad_t, [[6, 6], [14, 14]])
+
+    with backprop.GradientTape(persistent=True) as tape:
+      tape.watch(t)
+      one = matmul(t, b=t, transpose_a=True)
+      two = matmul(b=t, a=t, transpose_a=True)
+      three = matmul(a=t, b=t, transpose_a=True)
+
+    for output in [one, two, three]:
+      self.assertAllEqual(tape.gradient(output, t), [[6, 6], [14, 14]])
+
+  def testGradientInFunctionWithKeywordArguments(self):
+
+    @polymorphic_function.function
+    def f(x):
+      return backprop.gradients_function(lambda y: y * y, [0])(x)[0]
+
+    self.assertAllEqual(f(x=constant_op.constant(1.0)), 2.0)
+
+  def testFunctionHasNoSecondOrderGradient(self):
+
+    # This test needs nn_grad imported. We could just disable the lint error,
+    # but this way if the test is deleted we'll know the import isn't needed.
+    _ = nn_grad
+
+    v = variables.Variable(1.)
+
+    @polymorphic_function.function
+    def f(labels, logits):
+      return polymorphic_function.function(
+          nn_ops.sparse_softmax_cross_entropy_with_logits)(
+              labels=labels, logits=logits + v)
+
+    @polymorphic_function.function
+    def f_grad():
+      with backprop.GradientTape() as tape:
+        logits = constant_op.constant([1., 2.])
+        tape.watch(logits)
+        out = f(constant_op.constant(1), logits)
+      return tape.gradient(out, logits)
+    # Mainly we want to check that the function builds despite
+    # sparse_softmax_cross_entropy_with_logits not having a second-order
+    # gradient defined.
+    self.assertAllEqual([2], f_grad().shape)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/polymorphic_function/polymorphic_function_test_cpu_only.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class FunctionCpuOnlyTest(test.TestCase, parameterized.TestCase):
+  """Test that jit_compile=True correctly throws an exception if XLA is not available.
+
+  This test should only be run without `--config=cuda`, as that implicitly links
+  in XLA JIT.
+  """
+
+  def testJitCompileRaisesExceptionWhenXlaIsUnsupported(self):
+    if test.is_built_with_rocm() or test_util.is_xla_enabled():
+      return
+
+    with self.assertRaisesRegex(errors.UnimplementedError,
+                                'support for that platform linked in'):
+
+      @polymorphic_function.function(jit_compile=True)
+      def fn(x):
+        return x + x
+
+      fn([1, 1, 2, 3])
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_test.py

@@ -0,0 +1,45 @@
+# Copyright 2019 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.
+# ==============================================================================
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class FunctionTests(xla_test.XLATestCase):
+
+  def testVarInitializedInFunction(self):
+    with self.test_scope():
+      v_holder = []
+
+      @polymorphic_function.function
+      def add_var(x):
+        if not v_holder:
+          v = variables.Variable([1., 2.])
+          v_holder.append(v)
+          already_initialized = variables.Variable(3.)
+          with ops.init_scope():
+            already_initialized.assign(10.)
+          v_holder.append(already_initialized)
+        return v_holder[0] + v_holder[1] + x
+
+      self.assertAllClose([13., 14.], add_var(constant_op.constant(2.)))
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/polymorphic_function/saved_model_exported_concrete.py

@@ -0,0 +1,121 @@
+# 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.
+# ==============================================================================
+# pylint: disable=unidiomatic-typecheck
+"""ExportedConcreteFunction class and its associated functions.
+
+Part of saved model utils, a shim layer for working with
+functions exported/restored from saved models.
+This functionality should ultimately be moved into a first-class core API.
+"""
+
+import gc
+from tensorflow.python.eager.polymorphic_function import function_type_utils
+from tensorflow.python.trackable import base as trackable
+
+
+# TODO(kathywu): Delete this class when ConcreteFunctions can be copied with new
+# captures.
+class ExportedConcreteFunction(trackable.Trackable):
+  """A callable class that uses captures from the exported SavedModel graph."""
+  __slots__ = ("function", "tensor_map")
+
+  def __init__(self, function, tensor_map):
+    self.function = function
+    self.tensor_map = tensor_map
+
+  def __call__(self, *args, **kwargs):
+    bound_arguments = function_type_utils.canonicalize_function_inputs(
+        args, kwargs, self.function._function_type
+    )
+    filtered_flat_args = self.function._function_type.unpack_inputs(
+        bound_arguments
+    )
+    export_captures = _map_captures_to_created_tensors(
+        self.function.graph.captures, self.tensor_map, self.function)
+    return self.function._call_flat(filtered_flat_args, export_captures)
+
+
+def _map_captures_to_created_tensors(original_captures, tensor_map, function):
+  """Maps eager tensors captured by a function to Graph resources for export.
+
+  Args:
+    original_captures: A dictionary mapping from tensors captured by the
+      function to interior placeholders for those tensors (inside the function
+      body).
+    tensor_map: A dictionary mapping from resource tensors owned by the eager
+      context to resource tensors in the exported graph.
+    function: Function with the original captures. Only used when raising the
+      AssertionError.
+
+  Returns:
+    A list of stand-in tensors which belong to the exported graph, corresponding
+    to the function's captures.
+
+  Raises:
+    AssertionError: If the function references a resource which is not part of
+      `tensor_map`.
+  """
+  export_captures = []
+  for exterior, interior in original_captures:
+    mapped_resource = tensor_map.get(exterior, None)
+    if mapped_resource is None:
+      _raise_untracked_capture_error(function.name, exterior, interior)
+    export_captures.append(mapped_resource)
+  return export_captures
+
+
+def _raise_untracked_capture_error(function_name, capture,
+                                   internal_capture=None,
+                                   node_path=None):
+  """Raises AssertionError due to being unable to export a function."""
+  msg = ("Tried to export a function which references an 'untracked' resource. "
+         "TensorFlow objects (e.g. tf.Variable) captured by functions must be "
+         "'tracked' by assigning them to an attribute of a tracked object or "
+         "assigned to an attribute of the main object directly. See the "
+         "information below:"
+         f"\n\tFunction name = {function_name}")
+
+  if node_path is not None:
+    msg += f"\n\tPath to Function = {node_path}"
+
+  msg += f"\n\tCaptured Tensor = {capture}"
+  msg += f"\n\t{_get_trackable_parent_error_string(capture)}"
+
+  if internal_capture is not None:
+    msg += f"\n\tInternal Tensor = {internal_capture}"
+  raise AssertionError(msg)
+
+
+def _get_trackable_parent_error_string(capture):
+  """Gets error string with the capture's parent object."""
+  parent = getattr(capture, "_parent_trackable", None)
+  if parent is not None:
+    return f"Trackable referencing this tensor = {parent()}"
+
+  # Try to figure out where the resource came from by iterating over objects
+  # which reference it. This is slow and doesn't help us figure out how to
+  # match it to other objects when loading the SavedModel as a checkpoint,
+  # so we can't continue saving. But we can at least tell the user what
+  # needs attaching.
+  trackable_referrers = []
+  for primary_referrer in gc.get_referrers(capture):
+    if isinstance(primary_referrer, trackable.Trackable):
+      trackable_referrers.append(primary_referrer)
+    for secondary_referrer in gc.get_referrers(primary_referrer):
+      if isinstance(secondary_referrer, trackable.Trackable):
+        trackable_referrers.append(secondary_referrer)
+  return ("Trackable Python objects referring to this tensor "
+          "(from gc.get_referrers, limited to two hops) = [\n\t\t{}]"
+          .format("\n\t\t".join([repr(obj) for obj in trackable_referrers])))

tensorflow/python/eager/polymorphic_function/saved_model_utils.py

@@ -0,0 +1,82 @@
+# 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.
+# ==============================================================================
+# pylint: disable=unidiomatic-typecheck
+"""A shim layer for working with functions exported/restored from saved models.
+
+This functionality should ultimately be moved into a first-class core API.
+"""
+
+import numpy
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.saved_model import registration
+from tensorflow.python.trackable import base as trackable
+
+
+@registration.register_tf_serializable()
+class TrackableConstant(trackable.Trackable):
+  """Trackable class for captured constants."""
+  __slots__ = ("capture", "function", "_exported_tensor")
+
+  def __init__(self, capture, function):
+    self.capture = capture
+    self.function = function
+    self._exported_tensor = None
+
+  def _export_to_saved_model_graph(self, tensor_map, **unused_kwargs):
+    capture_constant_value = tensor_util.constant_value(self.capture)
+    if capture_constant_value is None:
+      raise ValueError(
+          f"Unable to save function {self.function.name} because it "
+          f"captures graph tensor {self.capture} from a parent function which "
+          "cannot be converted to a constant with `tf.get_static_value`.")
+
+    if numpy.prod(self.capture.shape.as_list()) > 1 and numpy.all(
+        capture_constant_value == capture_constant_value.flat[0]):
+      # For the common case of a constant array filled with the same
+      # value, rebuild the constant op specifically with the shape arg,
+      # since otherwise the whole array is written into the node def,
+      # causing performance and graph proto size issues (protos cannot be
+      # bigger than 2GB).
+      copied_tensor = constant_op.constant(
+          capture_constant_value.flat[0],
+          dtype=self.capture.dtype,
+          shape=self.capture.shape)
+    else:
+      copied_tensor = constant_op.constant(capture_constant_value)
+
+    tensor_map[self.capture] = copied_tensor
+    self._exported_tensor = copied_tensor
+    return [self.capture]
+
+  def _serialize_to_proto(self, object_proto=None, **kwargs):
+    object_proto.constant.operation = self._exported_tensor.op.name
+
+  @classmethod
+  def _deserialize_from_proto(cls, object_proto, operation_attributes,
+                              **kwargs):
+    tensor_proto = (
+        operation_attributes[object_proto.constant.operation]["value"].tensor)
+    ndarray = tensor_util.MakeNdarray(tensor_proto)
+    if dtypes.as_dtype(tensor_proto.dtype) == dtypes.string:
+      with ops.device("CPU"):
+        # String operations should be done on the CPU.
+        imported_constant = constant_op.constant(ndarray)
+    else:
+      imported_constant = constant_op.constant(ndarray)
+    return imported_constant

tensorflow/python/eager/polymorphic_function/tf_method_target.py

@@ -0,0 +1,51 @@
+# 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.
+# ==============================================================================
+"""Module for the TFMethodTarget Class."""
+
+import weakref
+
+from tensorflow.python.util import tf_inspect
+
+
+# When a method is bound to objects of this type, it allows AutoGraph to
+# recover a weak reference the original method's self pointer, so that it can
+# execute it consistent with class_method_to_instance_method's
+# bound_method_wrapper.
+# TODO(b/119246461): This is not pretty. Use a descriptor instead?
+class TfMethodTarget:
+  """Binding target for methods replaced by function and defun."""
+
+  __slots__ = ("weakrefself_target__", "weakrefself_func__")
+
+  def __init__(self, target, original_python_function):
+    self.weakrefself_target__ = target
+    self.weakrefself_func__ = weakref.ref(original_python_function)
+
+  @property
+  def target(self):
+    return self.weakrefself_target__()
+
+  @property
+  def target_class(self):
+    true_self = self.weakrefself_target__()
+    if tf_inspect.isclass(true_self):
+      # Class method
+      return true_self
+    else:
+      return true_self.__class__
+
+  def call(self, args, kwargs):
+    wrapped_fn = self.weakrefself_func__()
+    return wrapped_fn(self.weakrefself_target__(), *args, **kwargs)

tensorflow/python/eager/polymorphic_function/tracing_compilation.py

@@ -0,0 +1,379 @@
+# 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.
+# ==============================================================================
+"""Compile Python functions to TF graphs using tracing."""
+
+import contextlib
+import dataclasses
+import enum
+import threading
+from typing import Any, Callable, Dict, Optional, Tuple
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.capture import capture_container
+from tensorflow.core.function.polymorphism import function_cache as function_cache_lib
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.python.autograph.core import ag_ctx
+from tensorflow.python.eager import monitoring
+from tensorflow.python.eager.polymorphic_function import attributes as attributes_lib
+from tensorflow.python.eager.polymorphic_function import concrete_function as concrete_function_lib
+from tensorflow.python.eager.polymorphic_function import function_context
+from tensorflow.python.eager.polymorphic_function import function_type_utils
+from tensorflow.python.eager.polymorphic_function import transform
+from tensorflow.python.framework import func_graph as func_graph_module
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.profiler import trace
+from tensorflow.python.util import compat
+
+_graph_building_time_counter = monitoring.Counter(
+    "/tensorflow/core/tf_function/graph_building_time_usecs",
+    "Time for tf.function to build a graph (us).",
+)
+
+
+class ScopeType(enum.Enum):
+  """Enumerate scopes under which functions might be traced."""
+  NO_SCOPE = 1
+  VARIABLE_CREATION = 2
+  NO_VARIABLE_CREATION = 3
+
+
+@dataclasses.dataclass
+class TracingOptions:
+  """Configuration options for tracing."""
+  # Python function to trace.
+  python_function: Callable[[Any], Any] = lambda *args, **kwargs: None
+
+  # Name given to the traced function.
+  name: str = "function"
+
+  # Known FunctionType of the python function.
+  polymorphic_type: Optional[function_type_lib.FunctionType] = None
+
+  # Known default values for the python function parameters.
+  default_values: Optional[Dict[str, Any]] = None
+
+  # Identifies effecting scope under which the function is traced.
+  scope_type: ScopeType = ScopeType.NO_SCOPE
+
+  # FunctionDef attributes for traced function.
+  attributes: Optional[Dict[str, Any]] = None
+
+  # See https://www.tensorflow.org/guide/autograph for more information.
+  # If autograph is enabled.
+  autograph: bool = True
+  # Optional tuple of `tf.autograph.experimental.Feature` values.
+  autograph_options: Optional[Tuple[Any, ...]] = None
+
+  # Trace generalized functions where possible to avoid future retracing.
+  reduce_retracing: bool = False
+
+  # If true, graph of generated Function will be destroyed with the function.
+  bind_graph_to_function: bool = False
+
+  # A FunctionCache object that holds existing traced functions.
+  function_cache: Optional[function_cache_lib.FunctionCache] = None
+
+  # A FunctionCaptures object that tracks by-ref captures.
+  function_captures: Optional[capture_container.FunctionCaptures] = None
+
+  # If specified, guards tracing and function lookup
+  lock: Optional[threading.Lock] = None
+
+  def __post_init__(self):
+    if self.attributes:
+      for attribute in self.attributes:
+        if attribute not in attributes_lib.TRACING_COMPILATION_ALLOWLIST:
+          raise ValueError(
+              f"Tracing compilation does not support `{attribute}` as an"
+              " attribute."
+          )
+
+    if not self.polymorphic_type or self.default_values is None:
+      self.polymorphic_type = function_type_lib.FunctionType.from_callable(
+          self.python_function
+      )
+      self.default_values = function_type_lib.FunctionType.get_default_values(
+          self.python_function
+      )
+
+    self._input_signature = function_type_utils.to_input_signature(
+        self.polymorphic_type
+    )
+
+  @property
+  def is_pure(self):
+    return self.attributes and attributes_lib.IMPLEMENTS in self.attributes
+
+  @property
+  def input_signature(self):
+    return self._input_signature
+
+
+def call_function(args=None, kwargs=None, tracing_options=None):
+  """Traces a function for args and kwargs and calls it after."""
+  if not tracing_options:
+    tracing_options = TracingOptions()
+
+  args = args if args else ()
+  kwargs = kwargs if kwargs else {}
+  function = trace_function(
+      args=args, kwargs=kwargs, tracing_options=tracing_options
+  )
+
+  # Bind it ourselves to skip unnecessary canonicalization of default call.
+  bound_args = function.function_type.bind(*args, **kwargs)
+  flat_inputs = function.function_type.unpack_inputs(bound_args)
+  return function._call_flat(  # pylint: disable=protected-access
+      flat_inputs, captured_inputs=function.captured_inputs
+  )
+
+
+def trace_function(args=None, kwargs=None, tracing_options=None):
+  """Returns a `ConcreteFunction` specialized to inputs and execution context.
+
+  Compiles a Graph corresponding to the Python function logic and uses that
+  to generate a differentiable ConcreteFunction.
+
+  Args:
+    args: inputs to specialize on. Can be concrete values (e.g. 1) or
+      `tf.Tensor` or `tf.TensorSpec`.
+    kwargs: keyword inputs to specialize on. Concrete values (e.g. 1) or
+      `tf.Tensor` or `tf.TensorSpec`.
+    tracing_options: TracingOptions for the tracing process.
+  """
+  if not tracing_options:
+    tracing_options = TracingOptions()
+
+  args = args if args else ()
+  kwargs = kwargs if kwargs else {}
+
+  if tracing_options.input_signature and (args or kwargs):
+    # Check to see if a valid type can be generated from the args, kwargs
+    bound_args = function_type_utils.bind_function_inputs(
+        args,
+        kwargs,
+        tracing_options.polymorphic_type,
+        tracing_options.default_values,
+    )
+    args, kwargs = bound_args.args, bound_args.kwargs
+
+  with tracing_options.lock or contextlib.nullcontext():
+    if tracing_options.input_signature and not args and not kwargs:
+      args = tracing_options.input_signature
+      kwargs = {}
+
+    concrete_function = _maybe_define_function(
+        args, kwargs, tracing_options
+    )
+
+  if not tracing_options.bind_graph_to_function:
+    concrete_function._garbage_collector.release()  # pylint: disable=protected-access
+
+  return concrete_function
+
+
+def _maybe_define_function(args, kwargs, tracing_options):
+  """Gets a function for these inputs, defining it if necessary.
+
+  Args:
+    args: The varargs for the Python function.
+    kwargs: The keyword args for the Python function.
+    tracing_options: TracingOptions for the tracing process.
+
+  Returns:
+    A ConcreteFunction generated based on args, kwargs and tracing_options.
+
+  Raises:
+    ValueError: If inputs are incompatible with the input signature.
+    TypeError: If the function inputs include non-hashable objects
+    RuntimeError: If there's an internal bug (inconsistency) in handling
+      shape relaxation retracing.
+  """
+  bound_args = function_type_utils.canonicalize_function_inputs(
+      args,
+      kwargs,
+      tracing_options.polymorphic_type,
+      tracing_options.default_values,
+      tracing_options.is_pure,
+  )
+  args, kwargs = bound_args.args, bound_args.kwargs
+
+  if tracing_options.input_signature is not None:
+    args = (
+        *tracing_options.input_signature,
+        *args[len(tracing_options.input_signature) :],
+    )
+
+  current_func_context = function_context.make_function_context(
+      tracing_options.scope_type
+  )
+
+  capture_types = (
+      tracing_options.function_captures.capture_types
+      if tracing_options.function_captures
+      else {}
+  )
+  lookup_func_type, lookup_func_context = (
+      function_type_utils.make_canonicalized_monomorphic_type(
+          args,
+          kwargs,
+          capture_types,
+          tracing_options.polymorphic_type,
+      )
+  )
+
+  if tracing_options.function_cache is not None:
+    concrete_function = tracing_options.function_cache.lookup(
+        lookup_func_type, current_func_context
+    )
+  else:
+    concrete_function = None
+
+  if concrete_function is not None:
+    return concrete_function
+
+  # Use a timer for graph building only if not already inside a function. This
+  # avoids double counting graph building time for nested functions.
+  with monitoring.MonitoredTimer(
+      _graph_building_time_counter.get_cell()
+  ) if not ops.inside_function() else contextlib.nullcontext():
+    with trace.Trace("tf.function-graph_building"):
+      logging.vlog(
+          1,
+          "Creating new FuncGraph for Python function %r (key: %r, %r)",
+          tracing_options.python_function,
+          current_func_context,
+          lookup_func_type,
+      )
+      logging.vlog(
+          2, "Python function signature [args: %s] [kwargs: %s]", args, kwargs
+      )
+      ag_status = (
+          ag_ctx.Status.ENABLED
+          if tracing_options.autograph
+          else ag_ctx.Status.DISABLED
+      )
+      with ag_ctx.ControlStatusCtx(
+          status=ag_status, options=tracing_options.autograph_options
+      ):
+        func_graph = func_graph_module.FuncGraph(tracing_options.name)
+        if (
+            tracing_options.input_signature is None
+            and tracing_options.reduce_retracing
+            and tracing_options.function_cache
+        ):
+          target_func_type = tracing_options.function_cache.generalize(
+              current_func_context, lookup_func_type
+          )
+        else:
+          target_func_type = lookup_func_type
+        concrete_function = _create_concrete_function(
+            target_func_type, lookup_func_context, func_graph, tracing_options
+        )
+
+        if tracing_options.function_cache is not None:
+          tracing_options.function_cache.add(
+              concrete_function, current_func_context
+          )
+
+        return concrete_function
+
+
+def _create_concrete_function(
+    function_type, type_context, func_graph, tracing_options
+):
+  """Create a `ConcreteFunction` from `args`, `kwargs`, and `func_graph`."""
+  placeholder_context = trace_type.InternalPlaceholderContext(
+      func_graph, type_context.get_placeholder_mapping()
+  )
+  with func_graph.as_default():
+    placeholder_bound_args = function_type.placeholder_arguments(
+        placeholder_context
+    )
+
+  disable_acd = tracing_options.attributes and tracing_options.attributes.get(
+      attributes_lib.DISABLE_ACD, False
+  )
+  traced_func_graph = func_graph_module.func_graph_from_py_func(
+      tracing_options.name,
+      tracing_options.python_function,
+      placeholder_bound_args.args,
+      placeholder_bound_args.kwargs,
+      None,
+      func_graph=func_graph,
+      add_control_dependencies=not disable_acd,
+      arg_names=function_type_utils.to_arg_names(function_type),
+      create_placeholders=False,
+  )
+
+  transform.apply_func_graph_transforms(traced_func_graph)
+
+  graph_capture_container = traced_func_graph.function_captures
+
+  if tracing_options.function_captures:
+    # Maintain the list of all captures
+    tracing_options.function_captures.merge_by_ref_with(graph_capture_container)
+
+  # Create a new FunctionType including captures and outputs.
+  output_type = trace_type.from_value(
+      traced_func_graph.structured_outputs, type_context
+  )
+  traced_func_type = function_type_lib.FunctionType(
+      function_type.parameters.values(),
+      traced_func_graph.function_captures.capture_types,
+      return_annotation=output_type,
+  )
+
+  concrete_function = concrete_function_lib.ConcreteFunction.from_func_graph(
+      traced_func_graph,
+      traced_func_type,
+      tracing_options.attributes,
+      # Tell the ConcreteFunction to clean up its graph once it goes out of
+      # scope. This is not the default behavior since it gets used in some
+      # places (like Keras) where the FuncGraph lives longer than the
+      # ConcreteFunction.
+      shared_func_graph=False,
+  )
+  _set_arg_keywords(concrete_function)
+  transform.call_concrete_function_callbacks(concrete_function)
+
+  return concrete_function
+
+
+def _set_arg_keywords(concrete_function):
+  """Sets arg keywords for ConcreteFunction."""
+  seen_names = set()
+  concrete_function._arg_keywords = []  # pylint: disable=protected-access
+  prefix_counts = {}
+  graph = concrete_function.graph
+  num_captures = len(graph.internal_captures + graph.deferred_internal_captures)
+  num_positional = len(graph.inputs) - num_captures
+  for arg in concrete_function.graph.inputs[:num_positional]:
+    try:
+      user_arg_name = compat.as_str(arg.op.get_attr("_user_specified_name"))
+    except ValueError:
+      user_arg_name = "tensor_arg"
+    proposal = user_arg_name
+    while proposal in seen_names:
+      index = prefix_counts.get(user_arg_name, 1)
+      proposal = "{}_{}".format(user_arg_name, index)
+      prefix_counts[user_arg_name] = index + 1
+    seen_names.add(proposal)
+    concrete_function._arg_keywords.append(proposal)  # pylint: disable=protected-access
+  # Anything can be a positional argument, in the same order as .inputs
+  concrete_function._num_positional_args = (  # pylint: disable=protected-access
+      num_positional
+  )

tensorflow/python/eager/polymorphic_function/transform.py

@@ -0,0 +1,32 @@
+# 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.
+# ==============================================================================
+"""tf.function transformations implementation."""
+
+# TODO(fmuham): Move this logic to core/function when layered.
+# TODO(fmuham): Deprecate and migrate these as AtomicFunction transformations.
+FUNC_GRAPH_TRANSFORMS = []
+CONCRETE_FUNCTION_CALLBACKS = []
+
+
+def apply_func_graph_transforms(func_graph):
+  """Applies registered transformations to FuncGraph."""
+  for transform in FUNC_GRAPH_TRANSFORMS:
+    transform(func_graph)
+
+
+def call_concrete_function_callbacks(concrete_fn):
+  """Calls registered callbacks against new ConcreteFunctions."""
+  for callback in CONCRETE_FUNCTION_CALLBACKS:
+    callback(concrete_fn)

tensorflow/python/eager/pywrap_gradient_exclusions.cc

@@ -0,0 +1,928 @@
+/* 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.
+==============================================================================*/
+
+// Inputs/Outputs exclusion lists for GradientTape.
+//
+// This file is MACHINE GENERATED! Do not edit.
+// Generated by: tensorflow/python/eager/gen_gradient_input_output_exclusions.py
+
+#include "tensorflow/python/eager/pywrap_gradient_exclusions.h"
+
+#include "absl/types/optional.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/gtl/flatset.h"
+
+using tensorflow::string;
+
+namespace {
+// Keep static data in a format that's easy to init statically.
+struct OpIndexInfo {
+  const char *op_name;
+  int num_indices;
+  std::array<int, 4> unused_indices;
+};
+
+// Helper function to initialize FlatMap<string,FlatSet> from OpIndexInfo.
+template <typename T>
+auto OpGradientInfoInit(const T &a) {
+  auto *m = new tensorflow::gtl::FlatMap<string, tensorflow::gtl::FlatSet<int>>;
+  for (const auto &item : a) {
+    m->emplace(string(item.op_name),
+               tensorflow::gtl::FlatSet<int>(
+                   item.unused_indices.begin(),
+                   item.unused_indices.begin() + item.num_indices));
+  }
+  return m;
+}
+}  // namespace
+
+absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
+    const tensorflow::string &op_name) {
+  static std::array<OpIndexInfo, 367> a = {{
+      {"Acosh"},
+      {"AllToAll", 1, {0}},
+      {"ApproximateEqual"},
+      {"ArgMax"},
+      {"ArgMin"},
+      {"AsString"},
+      {"Asinh"},
+      {"Assign"},
+      {"AssignAdd"},
+      {"AssignSub"},
+      {"AudioSummary"},
+      {"AudioSummaryV2"},
+      {"AvgPool3DGrad", 1, {1}},
+      {"AvgPoolGrad", 1, {1}},
+      {"BatchNormWithGlobalNormalization", 1, {3}},
+      {"BatchToSpace", 1, {0}},
+      {"BatchToSpaceND", 1, {0}},
+      {"BiasAdd"},
+      {"BiasAddV1"},
+      {"BitwiseAnd"},
+      {"BitwiseOr"},
+      {"BitwiseXor"},
+      {"BroadcastGradientArgs"},
+      {"CSRSparseMatrixToSparseTensor"},
+      {"CTCBeamSearchDecoder"},
+      {"CTCGreedyDecoder"},
+      {"CTCLoss"},
+      {"CTCLossV2"},
+      {"Ceil"},
+      {"CheckNumerics"},
+      {"CheckNumericsV2"},
+      {"Cholesky"},
+      {"CollectivePermute", 1, {0}},
+      {"CompositeTensorVariantFromComponents"},
+      {"CompositeTensorVariantToComponents"},
+      {"Conj"},
+      {"ConjugateTranspose", 1, {0}},
+      {"Const"},
+      {"Conv2DBackpropFilter", 1, {1}},
+      {"Conv2DBackpropInput", 1, {0}},
+      {"Conv3DBackpropFilterV2", 1, {1}},
+      {"Conv3DBackpropInputV2", 1, {0}},
+      {"CropAndResize", 1, {3}},
+      {"CrossReplicaSum", 1, {0}},
+      {"Cumsum", 1, {0}},
+      {"DecodeBase64"},
+      {"DecodePaddedRaw"},
+      {"DecodeProtoV2"},
+      {"DecodeRaw"},
+      {"DeleteSessionTensor"},
+      {"DenseToCSRSparseMatrix"},
+      {"DenseToDenseSetOperation"},
+      {"DenseToSparseSetOperation"},
+      {"DepthToSpace"},
+      {"DepthwiseConv2dNativeBackpropFilter", 1, {1}},
+      {"DepthwiseConv2dNativeBackpropInput", 1, {0}},
+      {"Diag"},
+      {"DiagPart"},
+      {"DrawBoundingBoxes"},
+      {"EditDistance"},
+      {"Elu"},
+      {"EncodeBase64"},
+      {"EnsureShape"},
+      {"Enter"},
+      {"Equal"},
+      {"Erfinv"},
+      {"Exit"},
+      {"Exp"},
+      {"ExpandDims", 1, {1}},
+      {"ExtractGlimpse"},
+      {"FFT"},
+      {"FFT2D"},
+      {"FFT3D"},
+      {"FakeQuantWithMinMaxArgsGradient"},
+      {"FakeQuantWithMinMaxVarsGradient"},
+      {"FakeQuantWithMinMaxVarsPerChannelGradient"},
+      {"Fill"},
+      {"FixedLengthRecordReader"},
+      {"Floor"},
+      {"FloorDiv"},
+      {"FusedBatchNorm", 1, {2}},
+      {"FusedBatchNormGradV3", 1, {5}},
+      {"FusedBatchNormV2", 1, {2}},
+      {"FusedBatchNormV3", 1, {2}},
+      {"GenerateBoundingBoxProposals"},
+      {"GenerateVocabRemapping"},
+      {"GetSessionHandle"},
+      {"GetSessionHandleV2"},
+      {"GetSessionTensor"},
+      {"Greater"},
+      {"GreaterEqual"},
+      {"HSVToRGB"},
+      {"HashTable"},
+      {"HashTableV2"},
+      {"HistogramSummary"},
+      {"IFFT"},
+      {"IFFT2D"},
+      {"IFFT3D"},
+      {"Identity"},
+      {"IdentityN"},
+      {"IdentityReader"},
+      {"Imag"},
+      {"ImageProjectiveTransformV2", 1, {2}},
+      {"ImageProjectiveTransformV3", 2, {2, 3}},
+      {"ImageSummary"},
+      {"InitializeTable"},
+      {"InitializeTableFromTextFile"},
+      {"InitializeTableFromTextFileV2"},
+      {"InitializeTableV2"},
+      {"Inv"},
+      {"Invert"},
+      {"InvertPermutation"},
+      {"IsotonicRegression"},
+      {"LMDBReader"},
+      {"LeakyReluGrad", 1, {0}},
+      {"LeftShift"},
+      {"Less"},
+      {"LessEqual"},
+      {"LinSpace"},
+      {"LoadAndRemapMatrix"},
+      {"LogSoftmax"},
+      {"LogicalAnd"},
+      {"LogicalNot"},
+      {"LogicalOr"},
+      {"LookupTableFind"},
+      {"LookupTableFindV2"},
+      {"LookupTableInsert"},
+      {"LookupTableInsertV2"},
+      {"LookupTableSize"},
+      {"LookupTableSizeV2"},
+      {"LoopCond"},
+      {"MatrixBandPart", 1, {0}},
+      {"MatrixDiag"},
+      {"MatrixDiagPartV2", 1, {2}},
+      {"MatrixDiagPartV3", 1, {2}},
+      {"MatrixDiagV2", 4, {0, 2, 3, 4}},
+      {"MatrixDiagV3", 4, {0, 2, 3, 4}},
+      {"MatrixInverse"},
+      {"MatrixSetDiagV2", 1, {0}},
+      {"MatrixSetDiagV3", 1, {0}},
+      {"MatrixSolve", 1, {1}},
+      {"MatrixSquareRoot"},
+      {"MaxPool3DGrad", 1, {2}},
+      {"MaxPool3DGradGrad", 1, {2}},
+      {"MaxPoolGrad", 1, {2}},
+      {"MaxPoolGradGrad", 1, {2}},
+      {"MaxPoolGradV2", 1, {2}},
+      {"MirrorPad", 1, {0}},
+      {"MirrorPadGrad", 1, {0}},
+      {"Multinomial"},
+      {"MutableDenseHashTable"},
+      {"MutableDenseHashTableV2"},
+      {"MutableHashTable"},
+      {"MutableHashTableOfTensors"},
+      {"MutableHashTableOfTensorsV2"},
+      {"MutableHashTableV2"},
+      {"NcclAllReduce"},
+      {"NcclBroadcast"},
+      {"Ndtri"},
+      {"Neg"},
+      {"NextIteration"},
+      {"NonMaxSuppression"},
+      {"NonMaxSuppressionV2"},
+      {"NonMaxSuppressionWithOverlaps"},
+      {"NotEqual"},
+      {"NthElement", 1, {1}},
+      {"OneHot"},
+      {"OnesLike"},
+      {"OptionalGetValue"},
+      {"Pack"},
+      {"ParameterizedTruncatedNormal"},
+      {"ParseTensor"},
+      {"PlaceholderWithDefault"},
+      {"PopulationCount"},
+      {"PreventGradient"},
+      {"QuantizeAndDequantize"},
+      {"QuantizeAndDequantizeV2"},
+      {"QuantizeAndDequantizeV3"},
+      {"QuantizeAndDequantizeV4Grad", 1, {3}},
+      {"QueueClose"},
+      {"QueueDequeue"},
+      {"QueueDequeueMany"},
+      {"QueueDequeueUpTo"},
+      {"QueueSize"},
+      {"RaggedFillEmptyRows"},
+      {"RaggedRange"},
+      {"RandomCrop"},
+      {"RandomIndexShuffle"},
+      {"RandomShuffle"},
+      {"RandomStandardNormal"},
+      {"RandomUniform"},
+      {"Range"},
+      {"Rank"},
+      {"ReadVariableOp"},
+      {"ReaderNumRecordsProduced"},
+      {"ReaderNumWorkUnitsCompleted"},
+      {"ReaderRead"},
+      {"ReaderReadUpTo"},
+      {"ReaderReset"},
+      {"ReaderRestoreState"},
+      {"ReaderSerializeState"},
+      {"Real"},
+      {"Reciprocal"},
+      {"ReduceJoin"},
+      {"RefEnter"},
+      {"RefExit"},
+      {"RefIdentity"},
+      {"RefNextIteration"},
+      {"RegexReplace"},
+      {"Relu"},
+      {"Relu6"},
+      {"Relu6Grad", 1, {0}},
+      {"ReluGrad", 1, {0}},
+      {"Reshape", 1, {1}},
+      {"ResizeBicubic", 1, {1}},
+      {"ResizeBilinear", 1, {1}},
+      {"ResizeBilinearGrad"},
+      {"ResizeNearestNeighbor", 1, {1}},
+      {"Reverse", 1, {0}},
+      {"ReverseSequence", 1, {0}},
+      {"ReverseV2", 1, {0}},
+      {"RightShift"},
+      {"Rint"},
+      {"Roll", 1, {0}},
+      {"Round"},
+      {"Rsqrt"},
+      {"SampleDistortedBoundingBox"},
+      {"SampleDistortedBoundingBoxV2"},
+      {"ScalarSummary"},
+      {"ScaleAndTranslate", 1, {1}},
+      {"ScatterAdd"},
+      {"ScatterDiv"},
+      {"ScatterMul"},
+      {"ScatterNd", 2, {1, 2}},
+      {"ScatterNdAdd"},
+      {"ScatterNdNonAliasingAdd", 2, {0, 2}},
+      {"ScatterNdSub"},
+      {"ScatterNdUpdate"},
+      {"ScatterSub"},
+      {"SdcaFprint"},
+      {"SegmentSum", 1, {0}},
+      {"Select", 1, {2}},
+      {"Selu"},
+      {"SerializeTensor"},
+      {"SetSize"},
+      {"Shape"},
+      {"Sigmoid"},
+      {"Size"},
+      {"Slice", 1, {2}},
+      {"Softmax"},
+      {"SoftmaxCrossEntropyWithLogits", 1, {1}},
+      {"SpaceToBatch", 1, {0}},
+      {"SpaceToBatchND", 1, {0}},
+      {"SpaceToDepth"},
+      {"SparseAdd", 3, {2, 5, 6}},
+      {"SparseAddGrad"},
+      {"SparseDenseCwiseAdd"},
+      {"SparseFillEmptyRows"},
+      {"SparseMatrixMul"},
+      {"SparseMatrixNNZ"},
+      {"SparseMatrixSoftmax"},
+      {"SparseMatrixTranspose"},
+      {"SparseMatrixZeros"},
+      {"SparseReduceSum", 1, {1}},
+      {"SparseReorder", 1, {1}},
+      {"SparseSegmentMeanWithNumSegments", 1, {3}},
+      {"SparseSegmentSqrtNWithNumSegments", 1, {3}},
+      {"SparseSegmentSumWithNumSegments", 1, {3}},
+      {"SparseSlice", 2, {2, 4}},
+      {"SparseSoftmax", 1, {1}},
+      {"SparseSoftmaxCrossEntropyWithLogits", 1, {1}},
+      {"SparseSparseMaximum"},
+      {"SparseSparseMinimum"},
+      {"SparseTensorDenseAdd", 3, {1, 2, 3}},
+      {"SparseTensorToCSRSparseMatrix"},
+      {"SparseToSparseSetOperation"},
+      {"Split", 1, {1}},
+      {"Sqrt"},
+      {"SqrtGrad", 1, {1}},
+      {"Stack"},
+      {"StackClose"},
+      {"StackPop"},
+      {"StackPush"},
+      {"StatelessCase"},
+      {"StatelessMultinomial"},
+      {"StatelessParameterizedTruncatedNormal", 1, {1}},
+      {"StatelessRandomBinomial"},
+      {"StatelessRandomGammaV2", 1, {1}},
+      {"StatelessRandomGammaV3", 3, {1, 2, 3}},
+      {"StatelessRandomNormal"},
+      {"StatelessRandomNormalV2"},
+      {"StatelessRandomPoisson"},
+      {"StatelessRandomUniform"},
+      {"StatelessRandomUniformFullInt"},
+      {"StatelessRandomUniformFullIntV2"},
+      {"StatelessRandomUniformInt"},
+      {"StatelessRandomUniformIntV2"},
+      {"StatelessRandomUniformV2"},
+      {"StatelessTruncatedNormal"},
+      {"StatelessTruncatedNormalV2"},
+      {"StopGradient"},
+      {"StridedSliceGrad", 2, {0, 4}},
+      {"StringSplit"},
+      {"StringToHashBucket"},
+      {"StringToHashBucketFast"},
+      {"StringToHashBucketStrong"},
+      {"StringToNumber"},
+      {"TFRecordReader"},
+      {"Tanh"},
+      {"TensorArray"},
+      {"TensorArrayClose"},
+      {"TensorArrayCloseV2"},
+      {"TensorArrayCloseV3"},
+      {"TensorArrayGrad"},
+      {"TensorArrayGradV2"},
+      {"TensorArrayGradV3"},
+      {"TensorArrayGradWithShape"},
+      {"TensorArrayScatter", 2, {2, 3}},
+      {"TensorArrayScatterV2", 2, {2, 3}},
+      {"TensorArrayScatterV3", 2, {2, 3}},
+      {"TensorArraySize"},
+      {"TensorArraySizeV2"},
+      {"TensorArraySizeV3"},
+      {"TensorArraySplit", 3, {1, 2, 3}},
+      {"TensorArraySplitV2", 3, {1, 2, 3}},
+      {"TensorArraySplitV3", 3, {1, 2, 3}},
+      {"TensorArrayV2"},
+      {"TensorArrayV3"},
+      {"TensorArrayWrite", 2, {2, 3}},
+      {"TensorArrayWriteV2", 2, {2, 3}},
+      {"TensorArrayWriteV3", 2, {2, 3}},
+      {"TensorListConcatLists"},
+      {"TensorListConcatV2", 2, {1, 2}},
+      {"TensorListElementShape"},
+      {"TensorListFromTensor", 1, {1}},
+      {"TensorListGetItem", 1, {2}},
+      {"TensorListLength"},
+      {"TensorListPopBack"},
+      {"TensorListPushBack", 1, {0}},
+      {"TensorListPushBackBatch"},
+      {"TensorListScatter", 1, {2}},
+      {"TensorListScatterV2", 2, {2, 3}},
+      {"TensorListStack"},
+      {"TensorScatterAdd", 2, {0, 2}},
+      {"TensorScatterSub", 2, {0, 2}},
+      {"TensorScatterUpdate", 1, {0}},
+      {"TensorStridedSliceUpdate", 1, {0}},
+      {"TensorSummary"},
+      {"TensorSummaryV2"},
+      {"TextLineReader"},
+      {"Timestamp"},
+      {"TopKV2", 1, {1}},
+      {"Transpose", 1, {0}},
+      {"TridiagonalSolve", 1, {1}},
+      {"TruncateDiv"},
+      {"TruncatedNormal"},
+      {"Unpack"},
+      {"UnsortedSegmentSum", 2, {0, 2}},
+      {"VarIsInitializedOp"},
+      {"VariableShape"},
+      {"WholeFileReader"},
+      {"XlaClusterOutput"},
+      {"XlaSharding"},
+      {"XlaSpmdShardToFullShape"},
+      {"ZerosLike"},
+      {"_EagerConst"},
+      {"VarHandleOp"},
+  }};
+  static const auto &m = *OpGradientInfoInit(a);
+
+  auto it = m.find(op_name);
+  if (it != m.end()) {
+    return it->second;
+  }
+  return absl::nullopt;
+}
+
+absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
+    const tensorflow::string &op_name) {
+  static std::array<OpIndexInfo, 486> a = {{
+      {"Abs"},
+      {"AccumulateNV2"},
+      {"Acos"},
+      {"Add"},
+      {"AddN"},
+      {"AddV2"},
+      {"AllToAll"},
+      {"Angle"},
+      {"ApproxTopK", 1, {0}},
+      {"ApproximateEqual"},
+      {"ArgMax"},
+      {"ArgMin"},
+      {"AsString"},
+      {"Asin"},
+      {"Assert"},
+      {"Assign"},
+      {"AssignAdd"},
+      {"AssignSub"},
+      {"Atan"},
+      {"Atan2"},
+      {"Atanh"},
+      {"AudioSummary"},
+      {"AudioSummaryV2"},
+      {"AvgPool"},
+      {"AvgPool3D"},
+      {"AvgPool3DGrad"},
+      {"AvgPoolGrad"},
+      {"BatchMatMul"},
+      {"BatchMatMulV2"},
+      {"BatchMatMulV3"},
+      {"BatchNormWithGlobalNormalization"},
+      {"BatchToSpace"},
+      {"BatchToSpaceND"},
+      {"BesselI0"},
+      {"BesselJ0"},
+      {"BesselK0"},
+      {"BesselY0"},
+      {"Betainc"},
+      {"BiasAdd"},
+      {"BiasAddGrad"},
+      {"BiasAddV1"},
+      {"BitwiseAnd"},
+      {"BitwiseOr"},
+      {"BitwiseXor"},
+      {"BroadcastGradientArgs"},
+      {"BroadcastTo"},
+      {"CSRSparseMatrixToDense"},
+      {"CSRSparseMatrixToSparseTensor", 1, {1}},
+      {"CTCGreedyDecoder"},
+      {"CTCLoss", 1, {0}},
+      {"CTCLossV2", 1, {0}},
+      {"Cast"},
+      {"Ceil"},
+      {"CheckNumerics"},
+      {"CheckNumericsV2"},
+      {"ClipByValue"},
+      {"CollectivePermute"},
+      {"Complex"},
+      {"CompositeTensorVariantFromComponents"},
+      {"Concat"},
+      {"ConcatV2"},
+      {"Conj"},
+      {"ConjugateTranspose"},
+      {"Const"},
+      {"Conv2D"},
+      {"Conv2DBackpropFilter"},
+      {"Conv2DBackpropInput"},
+      {"Conv3D"},
+      {"Conv3DBackpropFilterV2"},
+      {"Conv3DBackpropInputV2"},
+      {"Cos"},
+      {"Cosh"},
+      {"CropAndResize"},
+      {"Cross"},
+      {"CrossReplicaSum"},
+      {"Cumprod"},
+      {"Cumsum"},
+      {"DecodeBase64"},
+      {"DecodePaddedRaw"},
+      {"DecodeRaw"},
+      {"DeleteSessionTensor"},
+      {"DenseToCSRSparseMatrix"},
+      {"DenseToDenseSetOperation"},
+      {"DenseToSparseSetOperation"},
+      {"DepthToSpace"},
+      {"DepthwiseConv2dNative"},
+      {"DepthwiseConv2dNativeBackpropFilter"},
+      {"DepthwiseConv2dNativeBackpropInput"},
+      {"Diag"},
+      {"DiagPart"},
+      {"Digamma"},
+      {"Dilation2D"},
+      {"Div"},
+      {"DivNoNan"},
+      {"DrawBoundingBoxes"},
+      {"DynamicPartition"},
+      {"EditDistance"},
+      {"Einsum"},
+      {"EluGrad"},
+      {"EncodeBase64"},
+      {"EncodeProto"},
+      {"EnsureShape"},
+      {"Enter"},
+      {"Equal"},
+      {"Erf"},
+      {"Erfc"},
+      {"Exit"},
+      {"ExpandDims"},
+      {"Expint"},
+      {"Expm1"},
+      {"ExtractGlimpse"},
+      {"FFT"},
+      {"FFT2D"},
+      {"FFT3D"},
+      {"FakeQuantWithMinMaxArgs"},
+      {"FakeQuantWithMinMaxArgsGradient"},
+      {"FakeQuantWithMinMaxVars"},
+      {"FakeQuantWithMinMaxVarsGradient"},
+      {"FakeQuantWithMinMaxVarsPerChannel"},
+      {"FakeQuantWithMinMaxVarsPerChannelGradient"},
+      {"Fill"},
+      {"FixedLengthRecordReader"},
+      {"Floor"},
+      {"FloorDiv"},
+      {"FloorMod"},
+      {"FractionalAvgPool", 1, {0}},
+      {"FresnelCos"},
+      {"FresnelSin"},
+      {"FusedBatchNorm", 3, {0, 1, 2}},
+      {"FusedBatchNormGrad"},
+      {"FusedBatchNormGradV2"},
+      {"FusedBatchNormGradV3"},
+      {"FusedBatchNormV2", 3, {0, 1, 2}},
+      {"FusedBatchNormV3", 3, {0, 1, 2}},
+      {"Gather"},
+      {"GatherNd"},
+      {"GatherV2"},
+      {"GenerateBoundingBoxProposals"},
+      {"GenerateVocabRemapping"},
+      {"GetSessionHandle"},
+      {"GetSessionHandleV2"},
+      {"GetSessionTensor"},
+      {"Greater"},
+      {"GreaterEqual"},
+      {"HSVToRGB"},
+      {"HashTable"},
+      {"HashTableV2"},
+      {"HistogramSummary"},
+      {"IFFT"},
+      {"IFFT2D"},
+      {"IFFT3D"},
+      {"IRFFT"},
+      {"IRFFT2D"},
+      {"Identity"},
+      {"IdentityN"},
+      {"IdentityReader"},
+      {"Igamma"},
+      {"Igammac"},
+      {"Imag"},
+      {"ImageProjectiveTransformV2"},
+      {"ImageProjectiveTransformV3"},
+      {"ImageSummary"},
+      {"InitializeTable"},
+      {"InitializeTableFromTextFile"},
+      {"InitializeTableFromTextFileV2"},
+      {"InitializeTableV2"},
+      {"InvGrad"},
+      {"Invert"},
+      {"InvertPermutation"},
+      {"IsotonicRegression", 1, {0}},
+      {"L2Loss"},
+      {"LMDBReader"},
+      {"LeakyRelu"},
+      {"LeakyReluGrad"},
+      {"LeftShift"},
+      {"Less"},
+      {"LessEqual"},
+      {"Lgamma"},
+      {"LinSpace"},
+      {"LoadAndRemapMatrix"},
+      {"Log"},
+      {"Log1p"},
+      {"LogMatrixDeterminant", 1, {0}},
+      {"LogicalAnd"},
+      {"LogicalNot"},
+      {"LogicalOr"},
+      {"LookupTableFind"},
+      {"LookupTableFindV2"},
+      {"LookupTableInsert"},
+      {"LookupTableInsertV2"},
+      {"LookupTableSize"},
+      {"LookupTableSizeV2"},
+      {"LoopCond"},
+      {"MatMul"},
+      {"MatrixBandPart"},
+      {"MatrixDiag"},
+      {"MatrixDiagPart"},
+      {"MatrixDiagPartV2"},
+      {"MatrixDiagPartV3"},
+      {"MatrixDiagV2"},
+      {"MatrixDiagV3"},
+      {"MatrixSetDiag"},
+      {"MatrixSetDiagV2"},
+      {"MatrixSetDiagV3"},
+      {"MaxPool3DGrad"},
+      {"MaxPool3DGradGrad"},
+      {"MaxPoolGrad"},
+      {"MaxPoolGradGrad"},
+      {"MaxPoolGradV2"},
+      {"MaxPoolWithArgmax", 1, {0}},
+      {"Maximum"},
+      {"Merge", 1, {0}},
+      {"MergeSummary"},
+      {"Minimum"},
+      {"MirrorPad"},
+      {"MirrorPadGrad"},
+      {"Mul"},
+      {"MulNoNan"},
+      {"Multinomial"},
+      {"MutableDenseHashTable"},
+      {"MutableDenseHashTableV2"},
+      {"MutableHashTable"},
+      {"MutableHashTableOfTensors"},
+      {"MutableHashTableOfTensorsV2"},
+      {"MutableHashTableV2"},
+      {"NcclAllReduce"},
+      {"NcclBroadcast"},
+      {"NcclReduce"},
+      {"Neg"},
+      {"NextAfter"},
+      {"NextIteration"},
+      {"NonMaxSuppression"},
+      {"NonMaxSuppressionV2"},
+      {"NonMaxSuppressionWithOverlaps"},
+      {"NotEqual"},
+      {"OneHot"},
+      {"OnesLike"},
+      {"OptionalFromValue"},
+      {"OptionalGetValue"},
+      {"Pack"},
+      {"Pad"},
+      {"PadV2"},
+      {"ParameterizedTruncatedNormal"},
+      {"ParseTensor"},
+      {"PlaceholderWithDefault"},
+      {"Polygamma"},
+      {"PopulationCount"},
+      {"PreventGradient"},
+      {"Print"},
+      {"Prod"},
+      {"QuantizeAndDequantize"},
+      {"QuantizeAndDequantizeV2"},
+      {"QuantizeAndDequantizeV3"},
+      {"QuantizeAndDequantizeV4"},
+      {"QuantizeAndDequantizeV4Grad"},
+      {"QueueClose"},
+      {"QueueEnqueue"},
+      {"QueueEnqueueMany"},
+      {"QueueSize"},
+      {"RFFT"},
+      {"RFFT2D"},
+      {"RaggedFillEmptyRows", 3, {0, 1, 2}},
+      {"RaggedGather"},
+      {"RaggedRange"},
+      {"RaggedTensorToSparse"},
+      {"RaggedTensorToTensor"},
+      {"RaggedTensorToVariant"},
+      {"RandomCrop"},
+      {"RandomIndexShuffle"},
+      {"RandomShuffle"},
+      {"RandomStandardNormal"},
+      {"RandomUniform"},
+      {"Range"},
+      {"Rank"},
+      {"ReadVariableOp"},
+      {"ReaderNumRecordsProduced"},
+      {"ReaderNumWorkUnitsCompleted"},
+      {"ReaderRead"},
+      {"ReaderReadUpTo"},
+      {"ReaderReset"},
+      {"ReaderRestoreState"},
+      {"ReaderSerializeState"},
+      {"Real"},
+      {"RealDiv"},
+      {"ReciprocalGrad"},
+      {"ReduceJoin"},
+      {"RefEnter"},
+      {"RefExit"},
+      {"RefIdentity"},
+      {"RefMerge", 1, {0}},
+      {"RefNextIteration"},
+      {"RefSwitch"},
+      {"RegexReplace"},
+      {"Relu6Grad"},
+      {"ReluGrad"},
+      {"Reshape"},
+      {"ResizeBicubic"},
+      {"ResizeBilinear"},
+      {"ResizeBilinearGrad"},
+      {"ResizeNearestNeighbor"},
+      {"ResourceGather"},
+      {"ResourceGatherNd"},
+      {"Reverse"},
+      {"ReverseSequence"},
+      {"ReverseV2"},
+      {"RightShift"},
+      {"Rint"},
+      {"Roll"},
+      {"Round"},
+      {"RsqrtGrad"},
+      {"SampleDistortedBoundingBox"},
+      {"SampleDistortedBoundingBoxV2"},
+      {"ScalarSummary"},
+      {"ScaleAndTranslate"},
+      {"ScatterAdd"},
+      {"ScatterDiv"},
+      {"ScatterMul"},
+      {"ScatterNd"},
+      {"ScatterNdAdd"},
+      {"ScatterNdNonAliasingAdd"},
+      {"ScatterNdSub"},
+      {"ScatterNdUpdate"},
+      {"ScatterSub"},
+      {"SdcaFprint"},
+      {"SdcaShrinkL1"},
+      {"SegmentMean"},
+      {"SegmentSum"},
+      {"Select"},
+      {"SeluGrad"},
+      {"SerializeTensor"},
+      {"SetSize"},
+      {"Shape"},
+      {"SigmoidGrad"},
+      {"Sign"},
+      {"Sin"},
+      {"Sinh"},
+      {"Size"},
+      {"SoftmaxCrossEntropyWithLogits", 1, {0}},
+      {"Softplus"},
+      {"SoftplusGrad"},
+      {"Softsign"},
+      {"SpaceToBatch"},
+      {"SpaceToBatchND"},
+      {"SpaceToDepth"},
+      {"SparseAdd", 2, {1, 2}},
+      {"SparseAddGrad"},
+      {"SparseConcat"},
+      {"SparseDenseCwiseAdd"},
+      {"SparseDenseCwiseDiv"},
+      {"SparseDenseCwiseMul"},
+      {"SparseFillEmptyRows", 3, {0, 1, 2}},
+      {"SparseMatMul"},
+      {"SparseMatrixAdd"},
+      {"SparseMatrixMatMul"},
+      {"SparseMatrixMul"},
+      {"SparseMatrixNNZ"},
+      {"SparseMatrixSparseMatMul"},
+      {"SparseMatrixTranspose"},
+      {"SparseMatrixZeros"},
+      {"SparseReduceSum"},
+      {"SparseReorder"},
+      {"SparseSegmentMean"},
+      {"SparseSegmentMeanWithNumSegments"},
+      {"SparseSegmentSqrtN"},
+      {"SparseSegmentSqrtNWithNumSegments"},
+      {"SparseSegmentSum"},
+      {"SparseSegmentSumWithNumSegments"},
+      {"SparseSlice", 2, {1, 2}},
+      {"SparseSoftmaxCrossEntropyWithLogits", 1, {0}},
+      {"SparseSparseMaximum"},
+      {"SparseSparseMinimum"},
+      {"SparseTensorDenseAdd"},
+      {"SparseTensorDenseMatMul"},
+      {"SparseTensorToCSRSparseMatrix"},
+      {"SparseToDense"},
+      {"SparseToSparseSetOperation"},
+      {"Spence"},
+      {"Split"},
+      {"SplitV"},
+      {"Square"},
+      {"SquaredDifference"},
+      {"Squeeze"},
+      {"Stack"},
+      {"StackClose"},
+      {"StackPop"},
+      {"StackPush"},
+      {"StatelessMultinomial"},
+      {"StatelessRandomBinomial"},
+      {"StatelessRandomNormal"},
+      {"StatelessRandomNormalV2"},
+      {"StatelessRandomPoisson"},
+      {"StatelessRandomUniform"},
+      {"StatelessRandomUniformFullInt"},
+      {"StatelessRandomUniformFullIntV2"},
+      {"StatelessRandomUniformInt"},
+      {"StatelessRandomUniformIntV2"},
+      {"StatelessRandomUniformV2"},
+      {"StatelessTruncatedNormal"},
+      {"StatelessTruncatedNormalV2"},
+      {"StopGradient"},
+      {"StridedSlice"},
+      {"StridedSliceGrad"},
+      {"StringJoin"},
+      {"StringSplit"},
+      {"StringToHashBucket"},
+      {"StringToHashBucketFast"},
+      {"StringToHashBucketStrong"},
+      {"StringToNumber"},
+      {"Sub"},
+      {"Sum"},
+      {"Switch"},
+      {"TFRecordReader"},
+      {"TPUEmbeddingActivations"},
+      {"TPUReplicatedInput"},
+      {"Tan"},
+      {"TanhGrad"},
+      {"TensorArray"},
+      {"TensorArrayClose"},
+      {"TensorArrayCloseV2"},
+      {"TensorArrayCloseV3"},
+      {"TensorArrayConcat", 1, {0}},
+      {"TensorArrayConcatV2", 1, {0}},
+      {"TensorArrayConcatV3", 1, {0}},
+      {"TensorArrayGather"},
+      {"TensorArrayGatherV2"},
+      {"TensorArrayGatherV3"},
+      {"TensorArrayGrad"},
+      {"TensorArrayGradV2"},
+      {"TensorArrayGradV3"},
+      {"TensorArrayGradWithShape"},
+      {"TensorArrayRead"},
+      {"TensorArrayReadV2"},
+      {"TensorArrayReadV3"},
+      {"TensorArraySize"},
+      {"TensorArraySizeV2"},
+      {"TensorArraySizeV3"},
+      {"TensorArrayV2"},
+      {"TensorArrayV3"},
+      {"TensorListConcat", 1, {0}},
+      {"TensorListConcatLists"},
+      {"TensorListConcatV2", 1, {0}},
+      {"TensorListElementShape"},
+      {"TensorListGather"},
+      {"TensorListGetItem"},
+      {"TensorListLength"},
+      {"TensorListPushBack"},
+      {"TensorListPushBackBatch"},
+      {"TensorListResize"},
+      {"TensorListScatter"},
+      {"TensorListScatterIntoExistingList"},
+      {"TensorListScatterV2"},
+      {"TensorListSetItem"},
+      {"TensorListSplit"},
+      {"TensorListStack"},
+      {"TensorScatterAdd"},
+      {"TensorScatterSub"},
+      {"TensorScatterUpdate"},
+      {"TensorStridedSliceUpdate"},
+      {"TensorSummary"},
+      {"TensorSummaryV2"},
+      {"TextLineReader"},
+      {"Tile"},
+      {"Timestamp"},
+      {"TopK", 1, {0}},
+      {"TopKV2", 1, {0}},
+      {"Transpose"},
+      {"TridiagonalMatMul"},
+      {"TruncateDiv"},
+      {"TruncatedNormal"},
+      {"Unpack"},
+      {"UnsortedSegmentSum"},
+      {"VarIsInitializedOp"},
+      {"VariableShape"},
+      {"WholeFileReader"},
+      {"Xdivy"},
+      {"XlaClusterOutput"},
+      {"XlaEinsum"},
+      {"XlaSharding"},
+      {"XlaSpmdFullToShardShape"},
+      {"XlaSpmdShardToFullShape"},
+      {"Xlog1py"},
+      {"Xlogy"},
+      {"ZerosLike"},
+      {"Zeta"},
+      {"_EagerConst"},
+      {"VarHandleOp"},
+  }};
+  static const auto &m = *OpGradientInfoInit(a);
+
+  auto it = m.find(op_name);
+  if (it != m.end()) {
+    return it->second;
+  }
+  return absl::nullopt;
+}

tensorflow/python/eager/pywrap_gradient_exclusions.h

@@ -0,0 +1,38 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_PYTHON_EAGER_PYWRAP_GRADIENT_EXCLUSIONS_H_
+#define TENSORFLOW_PYTHON_EAGER_PYWRAP_GRADIENT_EXCLUSIONS_H_
+
+#include "absl/types/optional.h"
+#include "tensorflow/core/lib/gtl/flatmap.h"
+#include "tensorflow/core/lib/gtl/flatset.h"
+
+// Lookup whether the Op with the given op_name has unused input indices.
+// Returns absl::nullopt if all inputs are used, set of unused indices
+// otherwise. Empty set indicates that all indices are unused. The latter is
+// necessary because sometimes it may not be possible to enumerate all indices
+// just using OpDef e.g. when there are `list(T)` or `N * T` type inputs.
+absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
+    const tensorflow::string& op_name);
+
+// Lookup whether the Op with the given op_name has unused output indices.
+// Returns absl::nullopt if all outputs are used, set of unused indices
+// otherwise. Empty set indicates that all indices are unused. The latter is
+// necessary because sometimes it may not be possible to enumerate all indices
+// just using OpDef e.g. when there are `list(T)` or `N * T` type outputs.
+absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
+    const tensorflow::string& op_name);
+
+#endif  // TENSORFLOW_PYTHON_EAGER_PYWRAP_GRADIENT_EXCLUSIONS_H_

tensorflow/python/eager/pywrap_tensor.h

@@ -0,0 +1,49 @@
+/* Copyright 2017 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.
+==============================================================================*/
+#ifndef TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_H_
+#define TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_H_
+
+// Must be included first
+// clang-format off
+#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
+// clang-format on
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/types.h"
+
+bool EagerTensor_CheckExact(const PyObject* o);
+int64_t PyEagerTensor_ID(const PyObject* tensor);
+tensorflow::DataType PyEagerTensor_Dtype(const PyObject* tensor);
+int64_t PyEagerTensor_NumElements(PyObject* tensor);
+TFE_TensorHandle* EagerTensor_Handle(const PyObject* o);
+
+namespace tensorflow {
+
+// Converts a value to a TFE_TensorHandle of a given dtype. The handle is
+// first allocated on CPU and then copied to a device identified by
+// device_name, unless it is nullptr.
+//
+// Note that an DT_INT32 handle is always kept on CPU regardless of the
+// device_name argument.
+TFE_TensorHandle* ConvertToEagerTensor(TFE_Context* ctx, PyObject* value,
+                                       DataType dtype,
+                                       const char* device_name = nullptr);
+
+PyObject* TFE_TensorHandleToNumpy(TFE_TensorHandle* handle, TF_Status* status);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_H_

tensorflow/python/eager/pywrap_tensor_conversion.cc

@@ -0,0 +1,85 @@
+/* Copyright 2019 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/python/eager/pywrap_tensor_conversion.h"
+
+#include "absl/log/check.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+
+namespace tensorflow {
+
+namespace {
+
+static monitoring::Counter<0>* const kScalarCacheHits =
+    monitoring::Counter<0>::New(
+        "/tensorflow/eager/python/scalar_cache_hits",
+        "Number of times a scalar TFE_TensorHandle was retrieved from cache");
+static monitoring::Counter<0>* const kScalarCacheMisses =
+    monitoring::Counter<0>::New(
+        "/tensorflow/eager/python/scalar_cache_misses",
+        "Number of times a scalar TFE_TensorHandle was not available in cache");
+
+}  // namespace
+
+TFE_TensorHandleCache* TFE_TensorHandleCache::Get() {
+  // TODO: b/169790439 - link with Context (in context.py) instead of having
+  // a static global?
+  static auto* cache = new TFE_TensorHandleCache();
+  return cache;
+}
+
+TFE_TensorHandle* TFE_TensorHandleCache::Lookup(
+    PyObject* value, DataType dtype, TFE_Context* ctx,
+    absl::string_view device_name) const {
+  CHECK(value != nullptr);  // Crash OK
+#ifdef Py_GIL_DISABLED
+  absl::MutexLock lock(mu_);
+#endif  // Py_GIL_DISABLED
+  const auto it =
+      cache_.find(LookupKey{PyObjectPtr{value}, dtype, ctx, device_name});
+  if (it == cache_.end()) {
+    kScalarCacheMisses->GetCell()->IncrementBy(1);
+    return nullptr;
+  }
+
+  kScalarCacheHits->GetCell()->IncrementBy(1);
+  TFE_TensorHandle* h = it->second;
+  unwrap(h)->Ref();
+  return h;
+}
+
+void TFE_TensorHandleCache::Insert(PyObject* value, DataType dtype,
+                                   TFE_Context* ctx,
+                                   absl::string_view device_name,
+                                   TFE_TensorHandle* h) {
+  CHECK(value != nullptr);  // Crash OK
+  CHECK(h != nullptr);      // Crash OK
+#ifdef Py_GIL_DISABLED
+  absl::MutexLock lock(mu_);
+#endif  // Py_GIL_DISABLED
+  auto [it, inserted] = cache_.try_emplace(
+      LookupKey{PyObjectPtr{value}, dtype, ctx, device_name}, h);
+  if (inserted) {
+    Py_INCREF(value);
+    unwrap(h)->Ref();
+  }
+}
+
+void TFE_TensorHandleCache::Clear() { DecrefUnrefAll(); }
+
+}  // namespace tensorflow

tensorflow/python/eager/pywrap_tensor_conversion.h

@@ -0,0 +1,151 @@
+/* Copyright 2019 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_CONVERSION_H_
+#define TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_CONVERSION_H_
+
+// Place `<locale>` before <Python.h> to avoid build failure in macOS.
+#include <locale>
+#include <string>
+#include <tuple>
+#include <utility>
+
+// The empty line above is on purpose as otherwise clang-format will
+// automatically move <Python.h> before <locale>.
+#include <Python.h>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/hash/hash.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+// Wrapper-class allowing to use Python hashing/comparison functions
+// for PyObject*.
+//
+// Note that unlike Safe_PyObjectPtr this class does not steal a
+// reference to a Python object. The caller is responsible for doing
+// Py_INCREF/Py_DECREF.
+class PyObjectPtr {
+ public:
+  template <typename H>
+  friend H AbslHashValue(H h, const PyObjectPtr& obj) {
+    Py_hash_t hash = PyObject_Hash(obj.ptr_);
+    CHECK_NE(hash, -1);        // Crash OK
+    CHECK(!PyErr_Occurred());  // Crash OK
+    return H::combine(std::move(h), hash);
+  }
+
+  explicit PyObjectPtr(PyObject* ptr) : ptr_(ptr) {}
+
+  explicit operator PyObject*() const { return ptr_; }
+
+  bool operator==(const PyObjectPtr& other) const {
+    // We require exact type equality to account for 0 == 0.0 == False.
+    if (Py_TYPE(ptr_) != Py_TYPE(other.ptr_)) {
+      return false;
+    }
+
+    bool result = PyObject_RichCompareBool(ptr_, other.ptr_, Py_EQ) > 0;
+    CHECK(!PyErr_Occurred());  // Crash OK
+    return result;
+  }
+
+ private:
+  PyObject* ptr_;
+};
+
+// Cache mapping PyObject* to the corresponding on-device TFE_TensorHandles.
+// Used to speed up ConvertToEagerTensor for scalars.
+// TODO: b/169790439 - move ConvertToEagerTensor here.
+class TFE_TensorHandleCache {
+ public:
+  static TFE_TensorHandleCache* Get();
+
+  TFE_TensorHandleCache() { cache_.reserve(64); }
+
+  TFE_TensorHandleCache(const TFE_TensorHandleCache&) = delete;
+  TFE_TensorHandleCache& operator=(const TFE_TensorHandleCache&) = delete;
+
+  ~TFE_TensorHandleCache() { DecrefUnrefAll(); }
+
+  TFE_TensorHandle* Lookup(PyObject* value, DataType dtype, TFE_Context* ctx,
+                           absl::string_view device_name) const;
+
+  void Insert(PyObject* value, DataType dtype, TFE_Context* ctx,
+              absl::string_view device_name, TFE_TensorHandle* h);
+
+  void Clear();
+
+ private:
+  // TODO: b/169790439 - Instead of `TFE_Context*` key, ideally Python's context
+  // object should have TFE_TensorHandleCache instance. Migrate once Python
+  // context object is backed by C++ data structure.
+  using Key = std::tuple<PyObjectPtr, DataType, TFE_Context*, std::string>;
+  using LookupKey =
+      std::tuple<PyObjectPtr, DataType, TFE_Context*, absl::string_view>;
+
+  struct KeyHash {
+    using is_transparent = void;
+
+    template <typename Tuple>
+    size_t operator()(const Tuple& t) const {
+      return absl::Hash<Tuple>{}(t);
+    }
+  };
+
+  struct KeyEqual {
+    using is_transparent = void;
+
+    template <typename Tuple1, typename Tuple2>
+    bool operator()(const Tuple1& lhs, const Tuple2& rhs) const {
+      return lhs == rhs;
+    }
+  };
+
+  auto ExtractCache() {
+#ifdef Py_GIL_DISABLED
+    absl::MutexLock lock(mu_);
+#endif  // Py_GIL_DISABLED
+    auto temp_cache = std::move(cache_);
+    cache_.clear();
+    return temp_cache;
+  }
+
+  void DecrefUnrefAll() {
+    auto temp_cache = ExtractCache();
+    for (const auto& [key, value] : temp_cache) {
+      Py_DECREF(static_cast<PyObject*>(std::get<0>(key)));
+      TFE_DeleteTensorHandle(value);
+    }
+  }
+
+#ifdef Py_GIL_DISABLED
+  mutable absl::Mutex mu_;
+#endif  // Py_GIL_DISABLED
+
+  // Under a GIL-enabled Python, guarded by the GIL. Under a no-GIL Python,
+  // guarded by mu_.
+  absl::flat_hash_map<Key, TFE_TensorHandle*, KeyHash, KeyEqual> cache_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_EAGER_PYWRAP_TENSOR_CONVERSION_H_

tensorflow/python/eager/pywrap_tensor_test.py

@@ -0,0 +1,105 @@
+# 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 TFE_TensorHandleToNumpy."""
+
+import numpy as np
+
+from tensorflow.python.eager import pywrap_tensor_test_util as util
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+
+
+class MyPythonObject:
+  pass
+
+
+def my_layer(x):
+  y = x**2
+  y.my_dynamic_attribute = MyPythonObject()
+  return y
+
+
+class PywrapTensorTest(test.TestCase):
+
+  def testGetScalarOne(self):
+    result = util.get_scalar_one()
+    self.assertIsInstance(result, np.ndarray)
+    self.assertAllEqual(result, 1.0)
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def test_no_leak(self):
+    x = constant_op.constant([1, 2, 3])
+    layer = my_layer(x)
+    for _ in range(int(1e2)):
+      layer = my_layer(x)
+    self.assertIsNotNone(layer)
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def test_no_leak_cycles(self):
+    for i in range(int(1e2)):
+      # use multiply to avoid cached tensors.
+      x = 1.0 * constant_op.constant([1.0, 1, 1, i])
+      y = 1.0 * constant_op.constant([1.0, 1, 2, i])
+      x.self_ref = lambda x: x
+      x.y = y
+      y.x = x
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def test_no_leak_shape(self):
+    for i in range(int(1e2)):
+      # use multiply to avoid cached tensors.
+      x = 1.0 * constant_op.constant([3.0, 1, 1, i])
+      x.shape.x = x
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def test_no_leak_handle_data(self):
+    for i in range(int(1e2)):
+      # use multiply to avoid cached tensors.
+      x = 1.0 * constant_op.constant([4.0, 1, 1, i])
+      x._handle_data = x
+
+  def test_delete_handle_data_raises(self):
+    x = constant_op.constant([1.0, 2.0])
+    with self.assertRaisesRegex(
+        AttributeError, "Cannot delete attribute '_handle_data'"
+    ):
+      del x._handle_data
+
+  def test_delete_tensor_shape_raises(self):
+    x = constant_op.constant([1.0, 2.0])
+    with self.assertRaisesRegex(
+        AttributeError, "Cannot delete attribute '_tensor_shape'"
+    ):
+      del x._tensor_shape
+
+  def test_handle_data_reentrancy(self):
+    x = constant_op.constant([1.0, 2.0])
+
+    class ReentrantObject:
+
+      def __init__(self, tensor):
+        self.tensor = tensor
+
+      def __del__(self):
+        self.tensor._handle_data = None
+
+    x._handle_data = ReentrantObject(x)
+    x._handle_data = None
+    self.assertIsNone(x._handle_data)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/pywrap_tensor_test_util.cc

@@ -0,0 +1,41 @@
+// 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/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/c/eager/c_api_test_util.h"
+#include "tensorflow/c/tf_status_helper.h"
+#include "tensorflow/python/eager/pywrap_tensor.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+using tensorflow::Pyo;
+using tensorflow::TF_StatusPtr;
+using tensorflow::TFE_TensorHandleToNumpy;
+
+PYBIND11_MODULE(pywrap_tensor_test_util, m) {
+  m.def("get_scalar_one", []() {
+    // Builds a TFE_TensorHandle and then converts to NumPy ndarray
+    // using TFE_TensorHandleToNumpy.
+    TFE_ContextOptions* opts = TFE_NewContextOptions();
+    TF_StatusPtr status(TF_NewStatus());
+    TFE_Context* ctx = TFE_NewContext(opts, status.get());
+    TFE_TensorHandle* handle = TestScalarTensorHandle(ctx, 1.0f);
+    auto result = Pyo(TFE_TensorHandleToNumpy(handle, status.get()));
+    TFE_DeleteTensorHandle(handle);
+    TFE_DeleteContext(ctx);
+    TFE_DeleteContextOptions(opts);
+    return result;
+  });
+}

tensorflow/python/eager/pywrap_tensor_test_util.pyi

@@ -0,0 +1,16 @@
+# 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 get_scalar_one() -> object: ...

tensorflow/python/eager/pywrap_tfe.h

@@ -0,0 +1,455 @@
+/* Copyright 2017 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_
+#define TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_
+
+// Place `<locale>` before <Python.h> to avoid build failure in macOS.
+#include <locale>
+
+// The empty line above is on purpose as otherwise clang-format will
+// automatically move <Python.h> before <locale>.
+#include <Python.h>
+
+#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+
+typedef absl::InlinedVector<TFE_TensorHandle*, 4UL> TFE_InputTensorHandles;
+typedef absl::InlinedVector<TFE_TensorHandle*, 2UL> TFE_OutputTensorHandles;
+
+// Execute a TensorFlow operation.
+//
+// 'device_name': Name of the device on which to execute the operation, or NULL
+//                for automatic selection.
+// 'op_name': Name of the TensorFlow op to execute.
+// 'inputs': An array of TFE_TensorHandle*'s of size 'num_inputs'. These tensors
+//           will be provided as input to the operation.
+// 'attrs': A Python tuple alternating names and attr values.
+// 'outputs': A pointer to a TFE_OutputTensorHandles in which outputs will
+//            placed. On success, its elements will be filled in and the
+//            caller takes ownership of each returned TFE_TensorHandle.
+//            'outputs' MUST be sized to be at least as large as the number
+//            of tensors produced by the operation and will be resized to
+//            the actual number of tensors produced.
+void TFE_Py_Execute(TFE_Context* ctx, const char* device_name,
+                    const char* op_name, TFE_InputTensorHandles* inputs,
+                    PyObject* attrs, TFE_OutputTensorHandles* outputs,
+                    TF_Status* out_status);
+
+// Execute a cancelable TensorFlow operation.
+//
+// Arguments as above (for TFE_Py_Execute), with the addition of:
+// 'cancellation_manager': A pointer to a TFE_CancellationManager that can be
+//                         used to cancel execution of the given operation.
+typedef struct TFE_CancellationManager TFE_CancellationManager;
+void TFE_Py_ExecuteCancelable(TFE_Context* ctx, const char* device_name,
+                              const char* op_name,
+                              TFE_InputTensorHandles* inputs, PyObject* attrs,
+                              TFE_CancellationManager* cancellation_manager,
+                              TFE_OutputTensorHandles* outputs,
+                              TF_Status* out_status);
+
+// Registers e as the Exception class for handling not ok Status. Returns
+// Py_None if registration succeeds, else throws a TypeError and returns NULL.
+//
+// This function is not thread-safe.
+PyObject* TFE_Py_RegisterExceptionClass(PyObject* e);
+
+// Registers e as the VSpace to use.
+// `vspace` must be a imperative_grad.py:VSpace named tuple.
+PyObject* TFE_Py_RegisterVSpace(PyObject* e);
+
+// Registers e as the Exception to be raised when the conditions of
+// TFE_Py_FastPathExecute_C have not been met. When this exception is set, it
+// is a signal to the calling code that it should fall back to the safer (and
+// more complete) code path.
+//
+// This function is not thread-safe.
+PyObject* TFE_Py_RegisterFallbackExceptionClass(PyObject* e);
+
+// Registers e as the gradient_function.
+// The registered function takes
+// (op_name, attrs, num_inputs, inputs, outputs, output_gradients) and returns
+// the input gradients. This function will not correctly be able to generate
+// gradients for functional ops - the gradients for those ops are calculated
+// through a different codepath (see function.py for additional information).
+//
+// This function is not thread-safe.
+PyObject* TFE_Py_RegisterGradientFunction(PyObject* e);
+
+// Registers e as the forward_gradient_function.  The registered function takes
+// (op_name, attrs, inputs, outputs, tangents) and returns the output
+// tangents. This function is used only for operations, not for custom gradients
+// or functional ops.
+//
+// This function is not thread-safe.
+PyObject* TFE_Py_RegisterJVPFunction(PyObject* e);
+
+namespace tensorflow {
+
+// Returns 0 if 'status' is TF_OK. Otherwise, raises an exception (using
+// `exception` if not nullptr, else using the class registered via
+// TFE_Py_RegisterExceptionClass), and returns -1.
+int MaybeRaiseExceptionFromTFStatus(TF_Status* status, PyObject* exception);
+
+}  // namespace tensorflow
+
+// Returns 0 if 'status' is ok. Otherwise, raises an exception (using
+// `exception` if not nullptr, else using the class registered via
+// TFE_Py_RegisterExceptionClass), and returns -1.
+int MaybeRaiseExceptionFromStatus(const absl::Status& status,
+                                  PyObject* exception);
+
+// Returns the string associated with the passed-in python object.
+const char* TFE_GetPythonString(PyObject* o);
+
+// Returns a unique id on each call.
+int64_t get_uid();
+
+// Wraps the output of get_uid as a Python Long object. Ownership is passed to
+// the caller.
+PyObject* TFE_Py_UID();
+
+// Deleter for Context objects, called from the Capsule that owns it.
+void TFE_DeleteContextCapsule(PyObject* context);
+
+// Returns true if o is an instance of EagerTensor, but not a subclass. Else
+// returns false.
+bool EagerTensor_CheckExact(const PyObject* o);
+
+// Helper function to construct a new EagerTensor from a TFE_TensorHandle.
+// This functions takes the ownership of the handle.
+PyObject* EagerTensorFromHandle(TFE_TensorHandle* handle,
+                                bool is_packed = false);
+
+// Extracts the handle inside EagerTensor object `o`. Returns nullptr on error.
+// This functions returns a unreferenced pointer to the handle.
+TFE_TensorHandle* EagerTensor_Handle(const PyObject* o);
+
+// Creates the `EagerTensor` class by subclassing `base_class` and returns the
+// newly created type, or nullptr on error.
+PyObject* TFE_Py_InitEagerTensor(PyObject* base_class);
+
+// Sets `profiler` as the current profiler to receive callbacks about events
+// on eager tensors. Currently, the only reported event is creation.
+// `profiler` is expected to have a `created(self, eager_tensor)` method that
+// takes the created tensor as its single argument.
+// Previous profiler, if any, is unset and will not receive any more
+// callbacks.
+// To unset the profiler, pass Py_None as the value of `profiler`.
+PyObject* TFE_Py_SetEagerTensorProfiler(PyObject* profiler);
+
+// Creates a new tape and adds it to the active set. `persistent` and
+// `watch_accessed_variables` must be `PyBool_Type` (`Py_True` or `Py_False`).
+PyObject* TFE_Py_TapeSetNew(PyObject* persistent,
+                            PyObject* watch_accessed_variables);
+
+// Removes the passed tape from the set of active tapes.
+void TFE_Py_TapeSetRemove(PyObject* tape);
+
+// Adds the passed tape to the set of active tapes.
+void TFE_Py_TapeSetAdd(PyObject* tape);
+
+// Returns true if the tape stack is empty.
+PyObject* TFE_Py_TapeSetIsEmpty();
+
+// Check if any backward tape should record an operation given inputs.
+//
+// Does not take forward accumulators into account.
+PyObject* TFE_Py_TapeSetShouldRecordBackprop(PyObject* tensors);
+
+// Determine possible gradient types, taking forward accumulators into account.
+//   - 0 if no tape will record (implies TFE_Py_TapeSetShouldRecordBackprop
+//     is false and no forward accumulator is watching)
+//   - 1 if first-order gradients may be requested
+//   - 2 if higher-order gradients may be requested
+PyObject* TFE_Py_TapeSetPossibleGradientTypes(PyObject* tensors);
+
+void TFE_Py_TapeWatch(PyObject* tape, PyObject* tensor);
+void TFE_Py_TapeSetDeleteTrace(int64_t tensor_id);
+
+// Stops any gradient recording on the current thread.
+//
+// Includes forward accumulators.
+void TFE_Py_TapeSetStopOnThread();
+
+// Restarts gradient recording on the current thread.
+void TFE_Py_TapeSetRestartOnThread();
+
+// Checks whether gradient recording is stopped on the current thread.
+PyObject* TFE_Py_TapeSetIsStopped();
+
+// Records an operation for the purpose of gradient computation.
+//
+// Arguments:
+//  - op_type is a string for the operation type, used in the backprop code
+//  - output_tensors are a list of Python Tensor objects output by the operation
+//  - input_tensors are a list of input Tensors to the recorded operation
+//  - backward_function is the function to be called during backprop or
+//    forwardprop to, given the gradients of the output tensors, produce the
+//    gradients of the input tensors. This function is automatically transposed
+//    during forwardprop.
+//  - forward_function is an optional special-case for forwardprop, taking input
+//    jvps and returning output jvps.
+//
+// Records an operation both for backprop (gradient tape) and forwardprop
+// (forward accumulator). Equivalent to calling both
+// TFE_Py_TapeSetRecordOperationBackprop and
+// TFE_Py_TapeSetRecordOperationForwardprop.
+PyObject* TFE_Py_TapeSetRecordOperation(PyObject* op_type,
+                                        PyObject* output_tensors,
+                                        PyObject* input_tensors,
+                                        PyObject* backward_function,
+                                        PyObject* forward_function);
+
+// Records an operation only for backprop (gradient tapes).
+//
+// Same arguments as TFE_Py_TapeSetRecordOperation.
+PyObject* TFE_Py_TapeSetRecordOperationBackprop(PyObject* op_type,
+                                                PyObject* output_tensors,
+                                                PyObject* input_tensors,
+                                                PyObject* backward_function);
+
+// Records an operation only for forwardprop (forward accumulators).
+//
+// Arguments:
+//  - op_type is a string for the operation type, used in the backprop code
+//  - output_tensors are a list of Python Tensor objects output by the operation
+//  - input_tensors are a list of input Tensors to the recorded operation
+//  - backward_function is the function to be called to, given the gradients of
+//    the output tensors, produce the gradients of the input tensors. This
+//    function is automatically transposed to produce output gradients given
+//    input gradients.
+//  - forwardprop_output_indices indicates any output_tensors which contain
+//    JVPs. Typically these will have come from TFE_Py_PackJVPs. May
+//    be None or an empty sequence if there are no JVP outputs from the
+//    operation.
+PyObject* TFE_Py_TapeSetRecordOperationForwardprop(
+    PyObject* op_type, PyObject* output_tensors, PyObject* input_tensors,
+    PyObject* backward_function, PyObject* forwardprop_output_indices);
+
+// Notifies all tapes that a variable has been accessed.
+void TFE_Py_TapeVariableAccessed(PyObject* variable);
+
+// Watches the given variable object on the given tape.
+void TFE_Py_TapeWatchVariable(PyObject* tape, PyObject* variable);
+
+// Computes a gradient based on information recorded on the tape.`tape` must
+// have been produced by TFE_Py_NewTape. `target` and `sources` must be python
+// lists of Tensor objects. `output_gradients` is either None or a python list
+// of either Tensor or None, and if not None should have the same length as
+// target.
+PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* target,
+                              PyObject* sources, PyObject* output_gradients,
+                              PyObject* sources_raw,
+                              PyObject* unconnected_gradients,
+                              TF_Status* status);
+
+// Execute a tensorflow operation assuming that all provided inputs are
+// correctly formatted (i.e. EagerTensors). If it doesn't find EagerTensors,
+// it will simply fail with a NotImplementedError.
+//
+// The "args" PyObject* is meant to be a tuple with the following structure:
+//  Item 1: The Python eager Context object
+//  Item 2: op_name: Name of the TensorFlow op to execute.
+//  Item 3: name: An optional name for the operation.
+//  Item 4 onwards: inputs - This is a list of inputs followed by a list of
+//        attrs. It is not necessary for type attrs to be present.
+//
+// Note: the device_name and op_callbacks, which were previously passed
+// as arguments, are now read via GetEagerContextThreadLocalData().
+//
+// This is named _C since there doesn't seem to be any way to make it visible
+// in the SWIG interface without renaming due to the use of the %native
+// directive.
+PyObject* TFE_Py_FastPathExecute_C(PyObject* args);
+
+// Record the gradient for a given op.
+PyObject* TFE_Py_RecordGradient(PyObject* op_name, PyObject* inputs,
+                                PyObject* attrs, PyObject* results,
+                                PyObject* forward_pass_name_scope);
+
+// Returns all variables watched by the given tape in the order those variables
+// were created.
+PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape);
+
+// Creates a new forward accumulator. Does not add it to the active set.
+PyObject* TFE_Py_ForwardAccumulatorNew(bool use_batch);
+
+// Adds a ForwardAccumulator to the active set, meaning it will watch executed
+// operations. It must not already be in the active set.
+PyObject* TFE_Py_ForwardAccumulatorSetAdd(PyObject* accumulator);
+// Removes a forward accumulator from the active set, meaning it will no longer
+// be watching operations.
+void TFE_Py_ForwardAccumulatorSetRemove(PyObject* accumulator);
+
+// Tell the forward accumulator `accumulator` to watch `tensor`, with a Tensor
+// tangent vector `tangent` of matching shape and dtype.
+void TFE_Py_ForwardAccumulatorWatch(PyObject* accumulator, PyObject* tensor,
+                                    PyObject* tangent);
+
+// Looks up the Jacobian-vector product of `tensor` in the forward accumulator
+// `accumulator`. Returns None if no JVP is available.
+PyObject* TFE_Py_ForwardAccumulatorJVP(PyObject* accumulator, PyObject* tensor);
+
+// Temporarily push or pop transient state for accumulators in the active set.
+//
+// Allows an accumulator which is currently processing an operation to
+// temporarily reset its state. This is useful when building forwardprop
+// versions of functions, where an accumulator will trigger function building
+// and then must process captured symbolic tensors while building it. Without
+// pushing and popping, accumulators ignore operations executed as a direct
+// result of their own jvp computations.
+PyObject* TFE_Py_ForwardAccumulatorPushState();
+PyObject* TFE_Py_ForwardAccumulatorPopState();
+
+// Collects state from all current forward accumulators related to `tensors`.
+//
+// This is useful for packing JVPs as function inputs before executing a
+// function which computes primals and JVPs at the same time.
+//
+// Does not include accumulators which are currently in the process of computing
+// a jvp (and so appear somewhere on the current execution stack) or any
+// accumulators more deeply nested.
+//
+// Includes JVPs for `tensors` and any higher-order JVPs for those
+// (recursively). Returns a two-element tuple (indices, jvps):
+//   indices: A sequence of sequences of two-element tuples. Each forward
+//       accumulator is represented as a sequence of tuples with (primal_index,
+//       jvp_index). Both integers index into the concatenated `tensors + jvps`
+//       array.
+//   jvps: A flat list of Tensors. Best interpreted as a sequence to be
+//       appended to `tensors`.
+PyObject* TFE_Py_PackJVPs(PyObject* tensors);
+
+// Variable Watcher methods.
+
+// Creates a new variable watcher and adds it to the set of active variable
+// watchers.
+PyObject* TFE_Py_VariableWatcherNew();
+
+// Removes the passed variable watcher from the set of active variable watchers.
+void TFE_Py_VariableWatcherRemove(PyObject* variable_watcher);
+
+// Notifies all variable watchers that a variable has been accessed.
+void TFE_Py_VariableWatcherVariableAccessed(PyObject* variable);
+
+// Returns all variables watched by the given variable_watcher in the order
+// those variables were created.
+PyObject* TFE_Py_VariableWatcherWatchedVariables(PyObject* variable_watcher);
+
+// Returns an EagerTensor of dimension [len(`tensors`)] containing
+// the `slice_dim`'th dimension of each tensor in `tensors`. In other words,
+// TFE_Py_TensorShapeSlice takes a slice of dimensions of tensors in
+// `tensors`. For example, if `tensors` contains tensors of with shapes
+// [1, 2, 3], [4, 5], [6, 7, 8, 9], TFE_Py_TensorShapeSlice called with
+// `slice_dim` equal to 1 will return [2, 5, 7].
+// On error, returns nullptr and sets python exception.
+// REQUIRES: `tensors` is a python list/tuple of EagerTensors
+// REQUIRES: `slice_dim` is non-negative and smaller than the rank of all
+//   tensors in `tensors`.
+PyObject* TFE_Py_TensorShapeSlice(PyObject* tensors, int slice_dim);
+
+// Returns the shape of this tensor's on-device representation.
+// The shape is represented as a Python tuple of integers.
+PyObject* TFE_Py_TensorShapeOnDevice(PyObject* tensor);
+
+void TFE_Py_EnableInteractivePythonLogging();
+
+// Sets the current Python eager Context object (defined
+// in eager/context.py). This function must be called at least once before
+// eager tensors are created.
+// If an error is encountered, sets python error and returns NULL. Else, returns
+// Py_None.
+//
+// Not thread-safe.
+// TODO(mdan): Retire this - non-Python users should only need the EagerContext.
+PyObject* TFE_Py_SetEagerContext(PyObject* py_context);
+
+// Returns the current eager Context object (defined in eager/context.py)
+// that was last set using TFE_Py_SetEagerContext.
+// If an error is encountered, sets python error and returns NULL.
+// The returned PyObject is "new", i.e. the caller must call Py_DECREF on it at
+// some point.
+PyObject* GetPyEagerContext();
+
+// These are exposed since there is SWIG code that calls these.
+// Returns a pre-allocated status if it exists.
+TF_Status* GetStatus();
+// Returns the pre-allocated status to the code.
+void ReturnStatus(TF_Status* status);
+
+namespace tensorflow {
+
+// Returns the DataType for the specified tensor.  Returns DT_INVALID if
+// PyObject is not a tensor.
+DataType PyTensor_DataType(PyObject* tensor);
+
+// Thread-local data associated with a Python eager Context object.
+//
+// TODO(edloper): Consider changing device_name and scope_name to a const char*
+// (with nullptr used for None). However, note that existing code (e.g.
+// TFE_TensorHandleCache::Lookup) assumes that the lifetime of these strings
+// extends beyond the point where their value is changed; so we'd need to make
+// sure that the strings stay alive (maybe using PyUnicode_InternInPlace?)
+struct EagerContextThreadLocalData {
+  bool is_eager = false;
+  bool invoking_op_callbacks = false;
+  tensorflow::Safe_PyObjectPtr device_name;
+  tensorflow::Safe_PyObjectPtr scope_name;
+  tensorflow::Safe_PyObjectPtr device_spec;
+  tensorflow::Safe_PyObjectPtr function_call_options;
+  tensorflow::Safe_PyObjectPtr executor;
+  tensorflow::Safe_PyObjectPtr op_callbacks;
+};
+
+// Create a thread-local-data structure associated with py_eager_context.
+// `is_eager` and `device_spec` are used to supply default values for those
+// fields whenever a new thread-local instance is created for py_eager_tensor.
+//
+// This function assumes that the Python GIL is held (and does not perform its
+// own locking).
+void MakeEagerContextThreadLocalData(PyObject* py_eager_context,
+                                     PyObject* is_eager,
+                                     PyObject* device_spec);
+
+// Returns the thread-local instance of EagerContextThreadLocalData that is
+// associated with the given Python Context object.  If an instance has not
+// yet been created for `py_eager_context` in this thread, then a new one is
+// created, and initialized with the default values specified in
+// MakeEagerContextThreadLocalData.
+EagerContextThreadLocalData* GetEagerContextThreadLocalData(
+    PyObject* py_eager_context);
+
+// Free data structures used to track py_eager_context.
+//
+// This frees global state associated with py_eager_context, as well as thread-
+// local state associated with py_eager_context and the current thread. If you
+// wish to destroy thread-local state associated with a single py_eager_context
+// for multiple threads, then you must call this method from each thread.
+//
+// Thread-local state associated with eager contexts is also automatically
+// cleaned up when the thread is destroyed.
+//
+// This function assumes that the Python GIL is held (and does not perform its
+// own locking).
+void DestroyEagerContextThreadLocalData(PyObject* py_eager_context);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_EAGER_PYWRAP_TFE_H_

tensorflow/python/eager/pywrap_tfe_test.py

@@ -0,0 +1,388 @@
+# Copyright 2018 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 low-level eager execution primitives."""
+
+import sys
+import traceback
+
+import numpy as np
+
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager import core
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import resource_variable_ops
+
+
+@test_util.with_eager_op_as_function
+class Tests(test.TestCase):
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_MatMulCorrectResponse(self):
+    a_2_by_2 = random_ops.random_uniform((2, 2))
+    b_2_by_2 = random_ops.random_uniform((2, 2))
+
+    a_100_by_784 = random_ops.random_uniform((100, 784))
+    b_100_by_784 = random_ops.random_uniform((100, 784))
+
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    self.assertAllClose(
+        math_ops.matmul(a_2_by_2, b_2_by_2),
+        pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None,
+                                          a_2_by_2, b_2_by_2, "transpose_a",
+                                          False, "transpose_b", False))
+    self.assertAllClose(
+        math_ops.matmul(a_100_by_784, b_100_by_784, transpose_b=True),
+        pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None,
+                                          a_100_by_784, b_100_by_784,
+                                          "transpose_a", False, "transpose_b",
+                                          True))
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_ResourceVariableMatMulCorrectResponse(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    a_2_by_2 = constant_op.constant(1.0, shape=[2, 2])
+    m = resource_variable_ops.ResourceVariable(a_2_by_2)
+    x = pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None, m,
+                                          m, "transpose_a", False,
+                                          "transpose_b", False)
+    y = pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None,
+                                          a_2_by_2, a_2_by_2, "transpose_a",
+                                          False, "transpose_b", False)
+
+    self.assertAllEqual(x, y)
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_TapeWrite(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    with backprop.GradientTape(persistent=True) as tape:
+      a_2_by_2 = constant_op.constant(1.0, shape=[2, 2])
+      tape.watch(a_2_by_2)
+      z = pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None,
+                                            a_2_by_2, a_2_by_2, "transpose_a",
+                                            False, "transpose_b", False)
+    dz_dy = tape.gradient(z, [a_2_by_2])[0]
+    self.assertAllEqual(dz_dy.numpy(),
+                        constant_op.constant(4.0, shape=[2, 2]).numpy())
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_ResourceVariableTapeWrite(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    with backprop.GradientTape(persistent=True) as tape:
+      a_2_by_2 = constant_op.constant(1.0, shape=[2, 2])
+      m = resource_variable_ops.ResourceVariable(a_2_by_2)
+      tape.watch(m)
+      z = pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None, m,
+                                            m, "transpose_a", False,
+                                            "transpose_b", False)
+    dz_dy = tape.gradient(z, [m])[0]
+    self.assertAllEqual(dz_dy.numpy(),
+                        constant_op.constant(4.0, shape=[2, 2]).numpy())
+
+  # Tests homogeneous list op
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_AddNCorrectResponse(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    a_2_by_2 = random_ops.random_uniform((2, 2))
+    b_2_by_2 = random_ops.random_uniform((2, 2))
+
+    self.assertAllClose(
+        math_ops.add_n([a_2_by_2, b_2_by_2]),
+        pywrap_tfe.TFE_Py_FastPathExecute(ctx, "AddN", None,
+                                          [a_2_by_2, b_2_by_2]))
+
+  # Tests homogeneous list op
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_AddNTapeWrite(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    a_2_by_2 = random_ops.random_uniform((2, 2))
+    b_2_by_2 = random_ops.random_uniform((2, 2))
+
+    with backprop.GradientTape(persistent=True) as tape:
+      tape.watch(a_2_by_2)
+      tape.watch(b_2_by_2)
+      z1 = pywrap_tfe.TFE_Py_FastPathExecute(ctx, "AddN", None,
+                                             [a_2_by_2, b_2_by_2])
+      z2 = math_ops.add_n([a_2_by_2, b_2_by_2])
+    dz1_dy = tape.gradient(z1, [a_2_by_2])[0]
+    dz2_dy = tape.gradient(z2, [a_2_by_2])[0]
+    self.assertAllEqual(dz1_dy.numpy(), dz2_dy.numpy())
+
+  # Tests heterogeneous list op
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_IdentityNCorrectResponse(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    a_2_by_2 = random_ops.random_uniform((2, 2))
+    b_2_by_2 = random_ops.random_uniform((2, 2))
+
+    self.assertAllClose(
+        array_ops.identity_n([a_2_by_2, b_2_by_2]),
+        pywrap_tfe.TFE_Py_FastPathExecute(ctx, "IdentityN", None,
+                                          [a_2_by_2, b_2_by_2]))
+
+  # Tests heterogeneous list op
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_IdentityNTapeWrite(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    a_2_by_2 = random_ops.random_uniform((2, 2))
+    b_2_by_2 = random_ops.random_uniform((2, 2))
+
+    with backprop.GradientTape(persistent=True) as tape:
+      tape.watch(a_2_by_2)
+      tape.watch(b_2_by_2)
+      z1 = pywrap_tfe.TFE_Py_FastPathExecute(ctx, "IdentityN",
+                                             None, [a_2_by_2, b_2_by_2])
+      z2 = array_ops.identity_n([a_2_by_2, b_2_by_2])
+    dz1_dy = tape.gradient(z1[0], [a_2_by_2])[0]
+    dz2_dy = tape.gradient(z2[0], [a_2_by_2])[0]
+    self.assertAllEqual(dz1_dy.numpy(), dz2_dy.numpy())
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastpathExecute_InvalidInputs(self):
+    a_2_by_2 = random_ops.random_uniform((2, 2))
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    assert ctx.executing_eagerly(
+    ), "The prototype doesn't contain C code for graph construction"
+    ctx_handle = ctx._handle  # pylint: disable=protected-access
+
+    # Not enough base params
+    with self.assertRaisesRegex(ValueError,
+                                "at least 3 items in the input tuple"):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, "Identity")
+
+    # Not enough inputs
+    with self.assertRaisesRegex(ValueError, "Expected to be at least 4, was 3"):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, "Identity", None)
+
+    # Bad type
+    with self.assertRaisesRegex(TypeError, "expected a string for op_name"):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, ctx_handle, None,
+                                        a_2_by_2)
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastPathExecute_InvalidAttributes(self):
+    split_dim = constant_op.constant(0, dtype=dtypes.int32)
+    value = constant_op.constant([0, 1, 2, 3], dtype=dtypes.float32)
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    with self.assertRaises(core._FallbackException):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, "Split", None,
+                                        split_dim, value, "num_split", -1)
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testFastPathExecute_VeryLargeOutputs(self):
+    split_dim = constant_op.constant(0, dtype=dtypes.int32)
+    value = constant_op.constant([0, 1, 2, 3], dtype=dtypes.float32)
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    with self.assertRaisesRegex(ValueError, "Number of outputs is too big"):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, "Split", None, split_dim, value,
+                                        "num_split", 1000000000000)
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testSlowPathExecute_VeryLargeOutputs(self):
+    split_dim = constant_op.constant(0, dtype=dtypes.int32)
+    value = [0, 1, 2, 3]
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    with self.assertRaises(core._FallbackException):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, "Split", None, split_dim, value,
+                                        "num_split", 1000000000000)
+
+    value = constant_op.constant(value)
+    attrs = ("num_split", 1000000000000, "T", value.dtype.as_datatype_enum)
+    with self.assertRaisesRegex(ValueError, "Number of outputs is too big"):
+      pywrap_tfe.TFE_Py_Execute(ctx._handle, None, "Split", [split_dim, value],
+                                attrs, 1000000000000)
+
+  @test_util.assert_no_new_tensors
+  @test_util.assert_no_garbage_created
+  def testInvalidNumOutputs(self):
+    with self.assertRaisesRegex(
+        Exception, r"Value for number_attr\(\) -1 < 0 \[Op:Split\]|"
+        r"Value for attr 'num_split' of -1 must be at least minimum 1"):
+      array_ops.split(value=[1, 2, 3], num_or_size_splits=-1)
+
+    with self.assertRaisesRegex(
+        Exception,
+        r"Value for attr 'num_split' of 0 must be at least minimum 1"):
+      array_ops.split(value=[1, 2, 3], num_or_size_splits=0)
+
+  def testIsFunction(self):
+    ctx = context.context()
+    self.assertFalse(ctx.has_function("not_a_function"))
+
+    @def_function.function
+    def f():
+      return 1.
+
+    self.assertTrue(ctx.has_function(f.get_concrete_function().name))
+
+  def testEagerExecute_InvalidType(self):
+    # Test case for GitHub issue 26879.
+    with ops.Graph().as_default():
+      a_2_by_2 = constant_op.constant(1.0, shape=[2, 2])
+      m = resource_variable_ops.ResourceVariable(a_2_by_2)
+      with self.assertRaisesRegex(TypeError,
+                                  "Expected list for 'values' argument"):
+        _ = array_ops_stack.stack(m, axis=1)
+
+  def testGraphResourceVariableRaisesFallback(self):
+    with ops.Graph().as_default():
+      a_2_by_2 = constant_op.constant(1.0, shape=[2, 2])
+      m = resource_variable_ops.ResourceVariable(a_2_by_2)
+    ctx = context.context()
+    ctx.ensure_initialized()
+    with self.assertRaises(core._FallbackException):
+      pywrap_tfe.TFE_Py_FastPathExecute(ctx, "MatMul", None, m, m,
+                                        "transpose_a", False, "transpose_b",
+                                        False)
+
+  def testOpDefDefaultType(self):
+    im = np.random.randint(  # pylint: disable=too-many-function-args
+        low=0,
+        high=65535,
+        size=100,
+        dtype=np.uint16).reshape(10, 10, 1)
+
+    context.ensure_initialized()
+
+    fastpath_dtype = test_ops.dtype_with_default_op(im).numpy()
+    slowpath_dtype = test_ops.dtype_with_default_op_eager_fallback(
+        im, None, context.context()).numpy()
+    # Ensure the fastpath and slowpath eager paths work.
+    self.assertEqual(fastpath_dtype, slowpath_dtype)
+
+    with ops.Graph().as_default(), self.cached_session():
+      graph_dtype_symbolic = test_ops.dtype_with_default_op(im)
+
+      graph_dtype = self.evaluate(graph_dtype_symbolic)
+    # Ensure the eager path matches the graph path.
+    self.assertEqual(fastpath_dtype, graph_dtype)
+
+    # Unfortunately, as of now, this doesn't work as expected on def_functions,
+    # since we convert the numpy arrays to tensors pre-tracing (which won't get
+    # overriddent by the default type).
+    @def_function.function
+    def func(im):
+      return test_ops.dtype_with_default_op(im)
+
+    function_dtype = func(im).numpy()
+    self.assertNotEqual(fastpath_dtype, function_dtype)
+
+    # Captures are OK, since they don't go through the conversion path.
+    @def_function.function
+    def func_captured():
+      return test_ops.dtype_with_default_op(im)
+
+    function_dtype = func_captured().numpy()
+    self.assertEqual(fastpath_dtype, function_dtype)
+
+  def testConvertFromArrayInterface(self):
+    context.ensure_initialized()
+    ctx = context.context()
+
+    class MyArrayClass(object):
+
+      def __init__(self):
+        self.array = np.random.random(16)
+
+      def __array__(self):
+        return self.array
+
+    a = MyArrayClass()
+    t = ops.EagerTensor(a, device=ctx.device_name, dtype=None)
+    self.assertAllEqual(t, a)
+
+    # TODO(b/147830189): Converting from EagerTensor should work.
+    # _ = ops.EagerTensor(t, device=ctx.device_name, dtype=None)
+
+    # TODO(b/147828820): Converting with tensors should work.
+    # _ = ops.EagerTensor([[t]], device=ctx.device_name, dtype=None)
+
+  def testFallbackErrorNotVisibleWhenFallbackMethodRaises(self):
+    ctx = context.context()
+    ctx.ensure_initialized()
+
+    try:
+      math_ops.mat_mul([[1., 1.] * 2], [[1., 1.] * 3])
+    except errors.InvalidArgumentError:
+      etype, value, tb = sys.exc_info()
+      full_exception_text = " ".join(
+          traceback.format_exception(etype, value, tb))
+
+    self.assertNotRegex(full_exception_text, "_FallbackException")
+
+  def testIntAttrThatDoesNotFitIn32Bits(self):
+    # Tests bug where int attributes >= 2**31 raised an exception on platforms
+    # where sizeof(long) = 32 bits.
+    ctx = context.context()
+    ctx.ensure_initialized()
+    shape = constant_op.constant([10])
+    minval = constant_op.constant(0)
+    maxval = constant_op.constant(10)
+    seed = 2**50
+    pywrap_tfe.TFE_Py_FastPathExecute(ctx, "RandomUniformInt", None,
+                                      shape, minval, maxval,
+                                      "seed", seed)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/record.py

@@ -0,0 +1,121 @@
+# 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.
+# ==============================================================================
+"""Gradient record utilities."""
+
+import contextlib
+
+from tensorflow.python import pywrap_tfe
+
+
+class VariableWatcher(object):
+  """A scope that tracks all trainable variable accesses within it.
+
+  This explicitly ignores variables that are not marked as trainable.
+
+  Sample usage:
+
+  var = tf.Variable(0.0)
+  with VariableWatcher() as variable_watcher:
+    var.assign_add(1.0)
+
+  assert variable_watcher.watched_variables == [var]
+  """
+
+  __slots__ = ["_variable_watcher"]
+
+  def __init__(self):
+    self._variable_watcher = None
+
+  def __enter__(self):
+    self._variable_watcher = pywrap_tfe.TFE_Py_VariableWatcherNew()
+    return self
+
+  def __exit__(self, typ, value, traceback):
+    pywrap_tfe.TFE_Py_VariableWatcherRemove(self._variable_watcher)
+
+  def watched_variables(self):
+    """Returns a tuple of variables accessed under this scope."""
+    return pywrap_tfe.TFE_Py_VariableWatcherWatchedVariables(
+        self._variable_watcher)
+
+
+@contextlib.contextmanager
+def stop_recording():
+  """Stop all gradient recording (backprop and forwardprop)."""
+  is_stopped = pywrap_tfe.TFE_Py_TapeSetIsStopped()
+  try:
+    if not is_stopped:
+      pywrap_tfe.TFE_Py_TapeSetStopOnThread()
+    yield
+  finally:
+    if not is_stopped:
+      pywrap_tfe.TFE_Py_TapeSetRestartOnThread()
+
+
+def should_record_backprop(tensors):
+  """Returns true if any tape in the stack watches any of these tensors.
+
+  Only takes GradientTapes into account, not forward accumulators.
+
+  Args:
+    tensors: Tensors to check, typically inputs to an operation.
+
+  Returns:
+    Boolean, whether any tape watches any of `tensors`.
+  """
+  return pywrap_tfe.TFE_Py_TapeSetShouldRecordBackprop(tensors)
+
+
+def record_operation(op_type, output_tensors, input_tensors, backward_function,
+                     forward_function=None):
+  """Records the operation on all tapes in the stack."""
+  pywrap_tfe.TFE_Py_TapeSetRecordOperation(op_type, output_tensors,
+                                           input_tensors, backward_function,
+                                           forward_function)
+
+
+def record_operation_backprop_only(op_type, output_tensors, input_tensors,
+                                   backward_function):
+  """Records the operation on all backward tapes in the stack."""
+  pywrap_tfe.TFE_Py_TapeSetRecordOperationBackprop(op_type, output_tensors,
+                                                   input_tensors,
+                                                   backward_function)
+
+
+def record_operation_forwardprop_only(op_type, output_tensors, input_tensors,
+                                      backward_function,
+                                      forwardprop_output_indices):
+  """Records the operation on all forward accumulators in the stack.
+
+  Args:
+    op_type: a string for the operation type, used in the backprop code
+    output_tensors: a list of Python Tensor objects output by the operation
+    input_tensors: a list of input Tensors to the recorded operation
+    backward_function: the function to be called to, given the gradients of the
+      output tensors, produce the gradients of the input tensors. This function
+      is automatically transposed to produce output gradients given input
+      gradients.
+    forwardprop_output_indices: indicates any output_tensors which contain JVPs.
+      Typically these will have come from TFE_Py_PackForwardGradients. May be
+      None or an empty sequence if there are no JVP outputs from the operation.
+  """
+  pywrap_tfe.TFE_Py_TapeSetRecordOperationForwardprop(
+      op_type, output_tensors, input_tensors, backward_function,
+      forwardprop_output_indices)
+
+
+def could_possibly_record():
+  """Returns True if any tape is active."""
+  return not pywrap_tfe.TFE_Py_TapeSetIsEmpty()

tensorflow/python/eager/record_test.py

@@ -0,0 +1,211 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Basic tests for gradients."""
+
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager import record
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import custom_gradient
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+# Importing nn_grad for the registration functions.
+from tensorflow.python.ops import nn_grad  # pylint: disable=unused-import
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import variables
+
+
+@custom_gradient.custom_gradient
+def two_outputs(a, b):
+  mm = math_ops.matmul(a, b)
+  r = math_ops.reduce_sum(mm)
+
+  def grad(dmm, dr):
+    return [
+        math_ops.matmul(dmm, b, transpose_b=True) +
+        math_ops.matmul(array_ops.ones_like(b * dr), b, transpose_b=True),
+        math_ops.matmul(a, dmm, transpose_b=True) +
+        math_ops.matmul(a, array_ops.ones_like(a) * dr, transpose_b=True)
+    ]
+
+  return [mm, r], grad
+
+
+@custom_gradient.custom_gradient
+def gradient_is_constant(x):
+  result = x * x
+
+  def grad(dr):
+    return [dr]
+
+  return result, grad
+
+
+class TapeTest(test.TestCase):
+
+  def testMultiOutput(self):
+
+    def fn(x, y):
+      c = x + y
+      # Multiple outputs from split.
+      d, f = array_ops.split(c, 2)
+      return d + f
+
+    a = constant_op.constant([[1., 0.], [0., 1.]])
+    b = constant_op.constant([[1., 2.], [3., 4.]])
+    da, db = backprop.gradients_function(fn, [0, 1])(a, b)
+    with context.graph_mode(), self.cached_session():
+      tf_a = constant_op.constant([[1, 0], [0, 1]], dtype=dtypes.float32)
+      tf_b = constant_op.constant([[1, 2], [3, 4]], dtype=dtypes.float32)
+      tf_c = tf_a + tf_b
+      tf_d, tf_f = array_ops.split(tf_c, 2, axis=1)
+      tf_e = tf_d + tf_f
+      tf_da, tf_db = gradients_impl.gradients(tf_e, [tf_a, tf_b])
+
+      self.assertAllEqual(da, self.evaluate(tf_da))
+      self.assertAllEqual(db, self.evaluate(tf_db))
+
+  def testBasicFunctional(self):
+
+    def forward(a, b):
+      mm = math_ops.matmul(a, b)
+      return math_ops.reduce_sum(mm)
+
+    aa = constant_op.constant([[1., 0.], [0., 1.]])
+    bb = constant_op.constant([[1., 2.], [3., 4.]])
+    da, = backprop.gradients_function(forward, ['a'])(aa, bb)
+    self.assertAllEqual(da,
+                        math_ops.matmul(
+                            array_ops.ones_like(aa),
+                            array_ops.transpose(bb)).numpy())
+
+  def testBasicFunctionalPositionalArg(self):
+
+    def forward(a, b):
+      mm = math_ops.matmul(a, b)
+      return math_ops.reduce_sum(mm)
+
+    aa = constant_op.constant([[1., 0.], [0., 1.]])
+    bb = constant_op.constant([[1., 2.], [3., 4.]])
+    da, = backprop.gradients_function(forward, [0])(aa, bb)
+    self.assertAllEqual(da,
+                        math_ops.matmul(
+                            array_ops.ones_like(aa),
+                            array_ops.transpose(bb)).numpy())
+
+  def testBasicFunctionalWithValue(self):
+
+    def forward(a, b):
+      mm = math_ops.matmul(a, b)
+      return math_ops.reduce_sum(mm)
+
+    aa = constant_op.constant([[1., 0.], [0., 1.]])
+    bb = constant_op.constant([[1., 2.], [3., 4.]])
+    val, (da,) = backprop.val_and_grad_function(forward, ['a'])(aa, bb)
+    self.assertAllEqual(da,
+                        math_ops.matmul(
+                            array_ops.ones_like(aa),
+                            array_ops.transpose(bb)))
+    self.assertAllEqual(val, forward(aa, bb))
+
+  def testTwoOutputs(self):
+
+    def fn(x, y):
+      mm, r = two_outputs(x, y)
+      return r + math_ops.reduce_sum(mm)
+
+    a = constant_op.constant([[1., 0.], [0., 1.]])
+    b = constant_op.constant([[1., 2.], [3., 4.]])
+    da, db = backprop.gradients_function(fn, [0, 1])(a, b)
+    with context.graph_mode(), self.cached_session():
+      tf_a = constant_op.constant([[1, 0], [0, 1]], dtype=dtypes.float32)
+      tf_b = constant_op.constant([[1, 2], [3, 4]], dtype=dtypes.float32)
+      tf_mm = math_ops.matmul(tf_a, tf_b)
+      tf_rr = 2 * math_ops.reduce_sum(tf_mm)
+      tf_da, tf_db = gradients_impl.gradients(tf_rr, [tf_a, tf_b])
+
+      self.assertAllEqual(da, self.evaluate(tf_da))
+      self.assertAllEqual(db, self.evaluate(tf_db))
+
+  def testGcTwoOutputs(self):
+
+    def fn(x, y):
+      return nn_ops.sparse_softmax_cross_entropy_with_logits(logits=x,
+                                                             labels=y)[0]
+
+    labels = constant_op.constant([0])
+    logits = constant_op.constant([[0.0]])
+    grad, = backprop.gradients_function(fn, [0])(logits, labels)
+    self.assertAllEqual(grad, [[0.0]])
+
+  def testTfTensor(self):
+
+    def fn(x):
+      return x
+
+    t = constant_op.constant(1.0)
+    g, = backprop.gradients_function(fn, [0])(t)
+    self.assertAllEqual(g, 1.0)
+
+
+class VariableWatcherTest(test.TestCase):
+
+  def testBasic(self):
+    var1 = variables.Variable(0.0)
+    var2 = variables.Variable(1.0)
+    with record.VariableWatcher() as variable_watcher:
+      var1.assign_add(1.0)
+      var2.assign_add(2.0)
+
+    self.assertAllEqual(variable_watcher.watched_variables(), (var1, var2))
+
+  def testNonTrainableVariables(self):
+    var1 = variables.Variable(0.0)
+    var2 = variables.Variable(1.0, trainable=False)
+    with record.VariableWatcher() as variable_watcher:
+      var1.assign_add(1.0)
+      var2.assign_add(2.0)
+
+    self.assertAllEqual(variable_watcher.watched_variables(), (var1,))
+
+  def testMultipleScopes(self):
+    var1 = variables.Variable(0.0)
+    var2 = variables.Variable(1.0)
+    with record.VariableWatcher() as variable_watcher1:
+      var1.assign_add(1.0)
+      with record.VariableWatcher() as variable_watcher2:
+        var2.assign_add(2.0)
+
+    # variable_watcher1 should see both vars and variable_watcher2 only sees
+    # var2
+    self.assertAllEqual(variable_watcher1.watched_variables(), (var1, var2))
+    self.assertAllEqual(variable_watcher2.watched_variables(), (var2,))
+
+  def testCreateVariables(self):
+    with record.VariableWatcher() as variable_watcher:
+      var1 = variables.Variable(0.0)
+      var2 = variables.Variable(1.0)
+      var1.assign_add(1.0)
+      var2.assign_add(2.0)
+
+    self.assertAllEqual(variable_watcher.watched_variables(), (var1, var2))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/remote.py

@@ -0,0 +1,279 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Helpers to connect to remote servers."""
+
+import copy
+
+from absl import logging
+
+from tensorflow.core.protobuf.tensorflow_server_pb2 import ServerDef
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import remote_utils
+from tensorflow.python.training import server_lib
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+_GRPC_PREFIX = "grpc://"
+_LOCAL_MASTERS = ("", "local")
+
+
+@tf_export("config.experimental_connect_to_host")
+def connect_to_remote_host(remote_host=None, job_name="worker"):
+  """Connects to a single machine to enable remote execution on it.
+
+  Will make devices on the remote host available to use. Note that calling this
+  more than once will work, but will invalidate any tensor handles on the old
+  remote devices.
+
+  Using the default job_name of worker, you can schedule ops to run remotely as
+  follows:
+  ```python
+  # When eager execution is enabled, connect to the remote host.
+  tf.config.experimental_connect_to_host("exampleaddr.com:9876")
+
+  with ops.device("job:worker/replica:0/task:1/device:CPU:0"):
+    # The following tensors should be resident on the remote device, and the op
+    # will also execute remotely.
+    x1 = array_ops.ones([2, 2])
+    x2 = array_ops.ones([2, 2])
+    y = math_ops.matmul(x1, x2)
+  ```
+
+  Args:
+    remote_host: a single or a list the remote server addr in host-port format.
+    job_name: The job name under which the new server will be accessible.
+
+  Raises:
+    ValueError: if remote_host is None.
+  """
+  if not remote_host:
+    raise ValueError("Must provide at least one remote_host")
+
+  remote_hosts = nest.flatten(remote_host)
+  cluster_spec = server_lib.ClusterSpec(
+      {job_name: [_strip_prefix(host, _GRPC_PREFIX) for host in remote_hosts]})
+
+  connect_to_cluster(cluster_spec)
+
+
+@tf_export("config.experimental_connect_to_cluster")
+def connect_to_cluster(cluster_spec_or_resolver,
+                       job_name="localhost",
+                       task_index=0,
+                       protocol=None,
+                       make_master_device_default=True,
+                       cluster_device_filters=None):
+  """Connects to the given cluster.
+
+  Will make devices on the cluster available to use. Note that calling this more
+  than once will work, but will invalidate any tensor handles on the old remote
+  devices.
+
+  If the given local job name is not present in the cluster specification, it
+  will be automatically added, using an unused port on the localhost.
+
+  Device filters can be specified to isolate groups of remote tasks to avoid
+  undesired accesses between workers. Workers accessing resources or launching
+  ops / functions on filtered remote devices will result in errors (unknown
+  devices). For any remote task, if no device filter is present, all cluster
+  devices will be visible; if any device filter is specified, it can only
+  see devices matching at least one filter. Devices on the task itself are
+  always visible. Device filters can be particially specified.
+
+  For example, for a cluster set up for parameter server training, the following
+  device filters might be specified:
+
+  ```python
+  cdf = tf.config.experimental.ClusterDeviceFilters()
+  # For any worker, only the devices on PS nodes and itself are visible
+  for i in range(num_workers):
+    cdf.set_device_filters('worker', i, ['/job:ps'])
+  # Similarly for any ps, only the devices on workers and itself are visible
+  for i in range(num_ps):
+    cdf.set_device_filters('ps', i, ['/job:worker'])
+
+  tf.config.experimental_connect_to_cluster(cluster_def,
+                                            cluster_device_filters=cdf)
+  ```
+
+  Args:
+    cluster_spec_or_resolver: A `ClusterSpec` or `ClusterResolver` describing
+      the cluster.
+    job_name: The name of the local job.
+    task_index: The local task index.
+    protocol: The communication protocol, such as `"grpc"`. If unspecified, will
+      use the default from `python/platform/remote_utils.py`.
+    make_master_device_default: If True and a cluster resolver is passed, will
+      automatically enter the master task device scope, which indicates the
+      master becomes the default device to run ops. It won't do anything if
+      a cluster spec is passed. Will throw an error if the caller is currently
+      already in some device scope.
+    cluster_device_filters: an instance of
+      `tf.train.experimental/ClusterDeviceFilters` that specify device filters
+      to the remote tasks in cluster.
+  """
+  if not context.executing_eagerly():
+    raise ValueError(
+        "`tf.config.experimental_connect_to_cluster` can only be called in "
+        "eager mode."
+    )
+  protocol = protocol or remote_utils.get_default_communication_protocol()
+  if isinstance(cluster_spec_or_resolver, server_lib.ClusterSpec):
+    cluster_spec = cluster_spec_or_resolver
+  elif isinstance(cluster_spec_or_resolver, cluster_resolver.ClusterResolver):
+    if cluster_spec_or_resolver.master() in _LOCAL_MASTERS:
+      # Do nothing if the master is local.
+      return
+    cluster_spec = cluster_spec_or_resolver.cluster_spec()
+  else:
+    raise ValueError(
+        "`cluster_spec_or_resolver` must be a `ClusterSpec` or a "
+        "`ClusterResolver`.")
+
+  cluster_def = copy.deepcopy(cluster_spec.as_cluster_def())
+  if cluster_device_filters:
+    if isinstance(cluster_device_filters, server_lib.ClusterDeviceFilters):
+      cluster_device_filters = copy.deepcopy(
+          cluster_device_filters._as_cluster_device_filters())  # pylint: disable=protected-access
+    else:
+      raise ValueError("`cluster_device_filters` must be an instance of "
+                       "`tf.train.experimental.ClusterDeviceFilters`.")
+
+  # Check whether the server def has changed. We need to do the check before the
+  # local job is added to the cluster.
+  is_server_def_changed = False
+  current_server_def = context.get_server_def()
+  if current_server_def and job_name not in cluster_spec.jobs:
+    for i, job in enumerate(current_server_def.cluster.job):
+      if job.name == job_name:
+        del current_server_def.cluster.job[i]
+  if (current_server_def is None or current_server_def.cluster != cluster_def or
+      current_server_def.job_name != job_name or
+      current_server_def.task_index != task_index):
+    is_server_def_changed = True
+
+  # Automatically add local job, if not part of the cluster spec.
+  if job_name not in cluster_spec.jobs:
+    local_port = pywrap_tfe.TF_PickUnusedPortOrDie()
+    job_def = cluster_def.job.add()
+    job_def.name = job_name
+    # TODO(fishx): Update this to make sure remote worker has valid ip address
+    # to connect with local.
+    job_def.tasks[0] = "localhost:{}".format(local_port)
+
+  if context.context().coordination_service is None:
+    service_type = remote_utils.coordination_service_type(protocol)
+    service_leader = ""
+    # Maybe enable coordination service for the communication protocol
+    # TODO(b/243839559): Fix UPTC + Coordination service crashing
+    # Check if cluster_spec_or_resolver is an instance of
+    #    tpu_cluster_resolver.TPUClusterResolver
+    if (isinstance(cluster_spec_or_resolver, cluster_resolver.ClusterResolver)
+        and hasattr(cluster_spec_or_resolver, "tpu_hardware_feature")):
+      service_leader = cluster_spec_or_resolver.get_coordination_service_leader(
+      )
+      # Maybe enable coordination service internally.
+      if cluster_spec_or_resolver.environment == "google":
+        is_uptc_sess = ".uptc-worker." in cluster_spec_or_resolver.master()
+        service_type = remote_utils.coordination_service_type(
+            protocol, is_uptc_sess)
+     # Enable coordination service for Cloud TPU.
+      else:
+        service_type = "standalone"
+
+    if service_type:
+      # If `enable_health_check` is true, coordination service agent would
+      # do connecting (and tasks would send heartbeat if connection is set up)
+      # while creating eager contexts. Enabling health check does not mutate
+      # coordination service.
+      context.context().configure_coordination_service(
+          service_type=service_type,
+          service_leader=service_leader,
+          enable_health_check=False)
+
+  default_session_config = copy.deepcopy(context.context().config)
+
+  for name in cluster_spec.jobs:
+    # assuming any of the non-local job is the worker jobs.
+    # should we use cluster_spec_or_resolver.get_job_name() instead when
+    # it is available?
+    # maybe consolicate this with the 'master' logic below
+    if name == job_name:
+      continue
+
+    default_session_config.experimental.collective_group_leader = (
+        f"/job:{name}/replica:0/task:0"
+    )
+
+  logging.info("default session config: %s", default_session_config)
+
+  server_def = ServerDef(
+      cluster=cluster_def,
+      job_name=job_name,
+      task_index=task_index,
+      protocol=protocol,
+      default_session_config=default_session_config,
+      cluster_device_filters=cluster_device_filters,
+  )
+
+  if is_server_def_changed:
+    context.set_server_def(server_def)
+  else:
+    context.update_server_def(server_def)
+
+  if make_master_device_default and isinstance(
+      cluster_spec_or_resolver,
+      cluster_resolver.ClusterResolver) and cluster_spec_or_resolver.master():
+    master = cluster_spec_or_resolver.master()
+    master_job_name = None
+    master_task_id = None
+    for job_name in cluster_spec.jobs:
+      for task_id in cluster_spec.task_indices(job_name):
+        task_address = cluster_spec.task_address(job_name, task_id)
+        if master in task_address or task_address in master:
+          master_job_name = job_name
+          master_task_id = task_id
+          break
+
+    if not master_job_name:
+      raise ValueError(
+          "`make_master_device_default` is set to True but cannot find "
+          "master %s in the cluster" % master)
+
+    master_device = "/job:{}/replica:0/task:{}".format(master_job_name,
+                                                       master_task_id)
+    master_device = device_util.canonicalize(master_device)
+    current_device = device_util.current()
+    if current_device:
+      current_device = device_util.canonicalize(current_device)
+    if current_device and current_device != master_device:
+      raise ValueError("`connect_to_cluster` is called inside existing device "
+                       "scope %s, which is different from the master device "
+                       "scope %s to enter. This is not allowed." %
+                       (current_device, master_device))
+    # TODO(b/138389076): Think of the entering device scope behavior in the
+    # failure recovery case when dealing with preemptions.
+    if not current_device:
+      logging.info("Entering into master device scope: %s", master_device)
+      ops.device(master_device).__enter__()
+
+
+def _strip_prefix(s, prefix):
+  return s[len(prefix):] if s.startswith(prefix) else s

tensorflow/python/eager/remote_benchmarks_test.py

@@ -0,0 +1,152 @@
+# Copyright 2019 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.
+# ==============================================================================
+r"""Benchmarks for remote worker eager execution.
+
+To run CPU benchmarks:
+  bazel run -c opt remote_benchmarks_test -- --benchmark_filter=.
+
+To run GPU benchmarks:
+  bazel run --config=cuda -c opt --copt="-mavx" remote_benchmarks_test -- \
+    --benchmark_filter=.
+"""
+
+import gc
+import time
+
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.training import server_lib
+
+
+def run_benchmark(func, num_iters, execution_mode=None):
+  ctx = context.context()
+  with context.execution_mode(execution_mode):
+    # call func to maybe warm up the GPU
+    func()
+    if execution_mode == context.ASYNC:
+      ctx.executor.wait()
+    start = time.time()
+    for _ in range(num_iters):
+      func()
+    if execution_mode == context.ASYNC:
+      ctx.executor.wait()
+    end = time.time()
+
+    return end - start
+
+
+class Foo(object):
+
+  def __init__(self, num_vars):
+    self._num_vars = num_vars
+    self._v = []
+
+  def __call__(self, inputs):
+    if not self._v:
+      for _ in range(self._num_vars):
+        self._v.append(variables.Variable(
+            random_ops.random_uniform([]), shape=[]))
+    for v in self._v:
+      inputs = inputs * v
+    return inputs
+
+
+class RemoteWorkerMicroBenchmarks(test.Benchmark):
+
+  def __init__(self):
+    # used for remote benchmarks
+    self._cached_server1 = server_lib.Server.create_local_server()
+    self._cached_server_target1 = self._cached_server1.target[len("grpc://"):]
+    self._cached_server2 = server_lib.Server.create_local_server()
+    self._cached_server_target2 = self._cached_server2.target[len("grpc://"):]
+
+  def _run(self, func, num_iters=1000, execution_mode=context.ASYNC):
+    total_time = run_benchmark(func, num_iters, execution_mode)
+    mean_us = total_time * 1e6 / num_iters
+    self.report_benchmark(
+        iters=num_iters,
+        wall_time=mean_us,
+        extras={"examples_per_sec": num_iters / total_time})
+
+  def benchmark_send(self):
+    remote.connect_to_remote_host(self._cached_server_target1)
+
+    x = random_ops.random_uniform((2, 2)).cpu()
+
+    @def_function.function
+    def remote_func(m):
+      return math_ops.matmul(m, m)
+
+    def func(m):
+      with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+        return remote_func(m)
+
+    self._run(lambda: func(x))
+    # NOTE(b/136184459): Force garbage collecting hanging resources before
+    # subsequent calls to set_server_def, to ensure the destroy resource ops are
+    # executed when their corresponding device and manager are still available.
+    gc.collect()
+
+  def benchmark_worker_recv(self):
+    remote.connect_to_remote_host(
+        [self._cached_server_target1, self._cached_server_target2])
+
+    with ops.device("job:worker/replica:0/task:1/device:CPU:0"):
+      v = variables.Variable(1.0)
+
+    @def_function.function
+    def remote_func():
+      return 1.0 + v
+
+    def func():
+      with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+        return remote_func()
+
+    self._run(func)
+    # NOTE(b/136184459): Force garbage collecting hanging resources before
+    # subsequent calls to set_server_def, to ensure the destroy resource ops are
+    # executed when their corresponding device and manager are still available.
+    gc.collect()
+
+  def benchmark_create_vars_inside_function(self):
+    remote.connect_to_remote_host(self._cached_server_target1)
+
+    def func():
+      with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+        layer = Foo(50)
+
+        @def_function.function
+        def remote_func():
+          with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+            return layer(random_ops.random_uniform([]))
+
+        return remote_func()
+
+    self._run(func, execution_mode=context.ASYNC, num_iters=100)
+    # NOTE(b/136184459): Force garbage collecting hanging resources before
+    # subsequent calls to set_server_def, to ensure the destroy resource ops are
+    # executed when their corresponding device and manager are still available.
+    gc.collect()
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/remote_cloud_tpu_test.py

@@ -0,0 +1,80 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Test that we can connect to a real Cloud TPU."""
+
+from absl import flags
+from absl.testing import absltest
+
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import remote
+from tensorflow.python.framework import config
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('tpu', '', 'Name of TPU to connect to.')
+flags.DEFINE_string('project', None, 'Name of GCP project with TPU.')
+flags.DEFINE_string('zone', None, 'Name of GCP zone with TPU.')
+
+flags.DEFINE_integer('num_tpu_devices', 8, 'The expected number of TPUs.')
+DEVICES_PER_TASK = 8
+
+EXPECTED_DEVICES_PRE_CONNECT = [
+    '/device:CPU:0',
+]
+EXPECTED_NEW_DEVICES_AFTER_CONNECT_TEMPLATES = [
+    '/job:worker/replica:0/task:{task}/device:CPU:0',
+    '/job:worker/replica:0/task:{task}/device:XLA_CPU:0',
+    '/job:worker/replica:0/task:{task}/device:TPU_SYSTEM:0',
+    '/job:worker/replica:0/task:{task}/device:TPU:0',
+    '/job:worker/replica:0/task:{task}/device:TPU:1',
+    '/job:worker/replica:0/task:{task}/device:TPU:2',
+    '/job:worker/replica:0/task:{task}/device:TPU:3',
+    '/job:worker/replica:0/task:{task}/device:TPU:4',
+    '/job:worker/replica:0/task:{task}/device:TPU:5',
+    '/job:worker/replica:0/task:{task}/device:TPU:6',
+    '/job:worker/replica:0/task:{task}/device:TPU:7',
+]
+
+
+class RemoteCloudTPUTest(absltest.TestCase):
+  """Test that we can connect to a real Cloud TPU."""
+
+  def test_connect(self):
+    # Log full diff on failure.
+    self.maxDiff = None  # pylint:disable=invalid-name
+
+    self.assertCountEqual(
+        EXPECTED_DEVICES_PRE_CONNECT,
+        [device.name for device in config.list_logical_devices()])
+
+    resolver = tpu_cluster_resolver.TPUClusterResolver(
+        tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project
+    )
+    remote.connect_to_cluster(resolver)
+
+    expected_devices = EXPECTED_DEVICES_PRE_CONNECT
+    for task in range(FLAGS.num_tpu_devices // DEVICES_PER_TASK):
+      expected_devices.extend([
+          template.format(task=task)
+          for template in EXPECTED_NEW_DEVICES_AFTER_CONNECT_TEMPLATES
+      ])
+
+    self.assertCountEqual(
+        expected_devices,
+        [device.name for device in config.list_logical_devices()])
+
+    tpu_cluster_resolver.initialize_tpu_system(resolver)
+
+if __name__ == '__main__':
+  absltest.main()

tensorflow/python/eager/remote_cluster_test.py

@@ -0,0 +1,692 @@
+# Copyright 2018 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 remote eager execution."""
+
+import os
+import threading
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.distributed_runtime.preemption import gen_check_preemption_op
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import executor
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import coordinator
+from tensorflow.python.training import server_lib
+
+JOB_NAME = "remote_device"
+
+
+def get_server_def(job_name, local_server_port, remote_server_addresses,
+                   task_index):
+  """Returns a server def with a single job + multiple tasks."""
+  cluster_def = cluster_pb2.ClusterDef()
+  job_def = cluster_def.job.add()
+  job_def.name = job_name
+  job_def.tasks[0] = "localhost:%d" % local_server_port
+
+  for i, remote_server_address in enumerate(remote_server_addresses, start=1):
+    job_def.tasks[i] = remote_server_address
+
+  server_def = tensorflow_server_pb2.ServerDef(
+      cluster=cluster_def,
+      job_name=job_name,
+      task_index=task_index,
+      protocol="grpc")
+  server_def.default_session_config.experimental.coordination_config.service_type = "standalone"
+
+  return server_def
+
+
+class DynamicClusterTest(test.TestCase, parameterized.TestCase):
+
+  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
+    super(DynamicClusterTest, self).__init__(methodName)
+    self._cached_server1 = server_lib.Server.create_local_server()
+    self._cached_server2 = server_lib.Server.create_local_server()
+    self._cached_server3 = server_lib.Server.create_local_server()
+    self._cached_server4 = server_lib.Server.create_local_server()
+
+    self._cached_server1_target = self._cached_server1.target[len("grpc://"):]
+    self._cached_server2_target = self._cached_server2.target[len("grpc://"):]
+    self._cached_server3_target = self._cached_server3.target[len("grpc://"):]
+    self._cached_server4_target = self._cached_server4.target[len("grpc://"):]
+
+    self.server_def_s1 = get_server_def(
+        JOB_NAME,
+        local_server_port=0,
+        remote_server_addresses=[self._cached_server1_target],
+        task_index=0)
+    self.server_def_s1_s2 = get_server_def(
+        JOB_NAME,
+        local_server_port=0,
+        remote_server_addresses=[
+            self._cached_server1_target, self._cached_server2_target
+        ],
+        task_index=0)
+    self.server_def_s1_s3 = get_server_def(
+        JOB_NAME,
+        local_server_port=0,
+        remote_server_addresses=[
+            self._cached_server1_target, self._cached_server3_target
+        ],
+        task_index=0)
+    self.server_def_s4_s3 = get_server_def(
+        JOB_NAME,
+        local_server_port=0,
+        remote_server_addresses=[
+            self._cached_server4_target, self._cached_server3_target
+        ],
+        task_index=0)
+    self.server_def_s1_s2_s3 = get_server_def(
+        JOB_NAME,
+        local_server_port=0,
+        remote_server_addresses=[
+            self._cached_server1_target, self._cached_server2_target,
+            self._cached_server3_target
+        ],
+        task_index=0)
+    self.server_def_s1_s2_s3_s4 = get_server_def(
+        JOB_NAME,
+        local_server_port=0,
+        remote_server_addresses=[
+            self._cached_server1_target, self._cached_server2_target,
+            self._cached_server3_target, self._cached_server4_target
+        ],
+        task_index=0)
+
+    self.device_local = "/job:%s/replica:0/task:0/device:CPU:0" % JOB_NAME
+    self.device_t1 = "/job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME
+    self.device_t2 = "/job:%s/replica:0/task:2/device:CPU:0" % JOB_NAME
+    self.device_t3 = "/job:%s/replica:0/task:3/device:CPU:0" % JOB_NAME
+    self.device_t4 = "/job:%s/replica:0/task:4/device:CPU:0" % JOB_NAME
+
+  def setUp(self):
+    super(DynamicClusterTest, self).setUp()
+    os.environ["TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE"] = str(False)
+    local_port = pywrap_tfe.TF_PickUnusedPortOrDie()
+    context.set_server_def(
+        server_def=get_server_def(
+            JOB_NAME,
+            local_server_port=local_port,
+            remote_server_addresses=[
+                self._cached_server1_target, self._cached_server2_target
+            ],
+            task_index=0))
+
+  def tearDown(self):
+    super(DynamicClusterTest, self).tearDown()
+    ops.device(None).__enter__()
+    context._reset_context()
+
+  def testCheckPreemption(self):
+    preemption_key = "TF_DEFAULT_PREEMPTION_NOTICE_KEY"
+    preemption_task = "/job:worker/task:0"
+
+    with ops.device(self.device_t1):
+      gen_check_preemption_op.check_preemption(preemption_key=preemption_key)
+
+    # Simulate a preemption notifier callback invocation.
+    context.context().set_config_key_value(preemption_key, preemption_task)
+    with self.assertRaises(errors.AbortedError) as cm:
+      with ops.device(self.device_t2):
+        gen_check_preemption_op.check_preemption(preemption_key=preemption_key)
+    self.assertEqual(
+        cm.exception.experimental_payloads.get(
+            b"type.googleapis.com/tensorflow.distributed_runtime.WorkerPreemption"
+        ), preemption_task.encode())
+
+  @test_util.run_in_async_and_sync_mode
+  def testServerAdded(self):
+    """Add a server to cluster, and run remote ops on it."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    context.update_server_def(server_def=self.server_def_s1_s2_s3)
+    with ops.device(self.device_t3):
+      x2 = array_ops.ones([2, 2])
+
+    # Test new server accessing resources on old server
+    with ops.device(self.device_t3):
+      y = math_ops.matmul(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # Test old server accessing resources on new server
+    with ops.device(self.device_t2):
+      y = math_ops.matmul(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testServerRemoved(self):
+    """Remove a server from cluster, and run ops on cluster."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+    with ops.device(self.device_t2):
+      x2 = array_ops.ones([2, 2])
+
+    with ops.device(self.device_t1):
+      y = math_ops.matmul(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    context.update_server_def(server_def=self.server_def_s1)
+    with ops.device(self.device_t1):
+      y = math_ops.matmul(x1, x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # Running ops on removed server s2 throws an exception
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      with ops.device(self.device_t2):
+        y = math_ops.matmul(x1, x2)
+    self.assertIn("unknown device", cm.exception.message)
+
+    # TODO(haoyuzhang): raise and catch exception when accessing tensors on
+    # the removed servers.
+
+  @test_util.run_in_async_and_sync_mode
+  def testServerReplaced(self):
+    """Replace remote host_port for a task, and run ops on cluster."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    context.update_server_def(server_def=self.server_def_s1_s3)
+    with ops.device(self.device_t2):
+      y = math_ops.matmul(x1, x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionServerAdded(self):
+    """Add a server to cluster, and run remote function on it."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    context.update_server_def(server_def=self.server_def_s1_s2_s3)
+    with ops.device(self.device_t3):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    with ops.device(self.device_t3):
+      x2 = array_ops.ones([2, 2])
+    with ops.device(self.device_t1):
+      y = worker_fn(x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionRunsAtMostOnceInvokedOnce(self):
+    """Run a function decorated with `function_runs_at_most_once`."""
+    @def_function.function(
+        experimental_attributes={"function_runs_at_most_once": True}
+    )
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    x = array_ops.ones([2, 2])
+    y = worker_fn(x)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+    # The kernel will be destroyed at this point. The purpose of this test is
+    # to ensure it tears down cleanly because we have already released the
+    # function handle.
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionRunsAtMostOnceInvokedTwice(self):
+    """Run a function decorated with `function_runs_at_most_once` twice."""
+    @def_function.function(
+        experimental_attributes={"function_runs_at_most_once": True}
+    )
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    x = array_ops.ones([2, 2])
+    y = worker_fn(x)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # Since the worker_fn can run at most once a second invocation should fail.
+    with self.assertRaisesRegex(errors.NotFoundError, "Handle.* not found"):
+      worker_fn(x)
+      if context.is_async():
+        context.async_wait()
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionServerRemoved(self):
+    """Remove a server from cluster, and run ops on cluster."""
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    context.update_server_def(server_def=self.server_def_s1)
+
+    with ops.device(self.device_t1):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # Running functions on removed server s2 throws an exception
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      with ops.device(self.device_t2):
+        y = worker_fn(x1)
+    self.assertIn(" unknown device", cm.exception.message)
+
+    # TODO(haoyuzhang): raise and catch exception when accessing tensors on
+    # the removed servers.
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionServerRemovedAddedBack(self):
+    """Add and remove a server, and run functions on cluster."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    context.update_server_def(server_def=self.server_def_s1_s2_s3)
+    with ops.device(self.device_t3):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    context.update_server_def(server_def=self.server_def_s1_s2)
+    with ops.device(self.device_t2):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    context.update_server_def(server_def=self.server_def_s1_s2_s3)
+    with ops.device(self.device_t3):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionServerReplaced(self):
+    """Replace remote host_port for a task, and run functions on cluster."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    context.update_server_def(server_def=self.server_def_s1_s3)
+    with ops.device(self.device_t2):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testFunctionRegisteredAndRemoved(self):
+    """Update cluster when other function are registered and removed."""
+    with ops.device(self.device_local):
+      x1 = array_ops.ones([2, 2])
+
+    num_calls = 30
+    self._coord = coordinator.Coordinator()
+
+    def update_server_def_fn():
+      with self._coord.stop_on_exception():
+        for i in range(num_calls):
+          context.update_server_def(
+              server_def=(self.server_def_s1_s2 if i %
+                          2 == 0 else self.server_def_s1_s3))
+
+    t = threading.Thread(target=update_server_def_fn)
+    t.start()
+
+    for _ in range(num_calls):
+
+      @def_function.function
+      def worker_fn(i):
+        return math_ops.matmul(i, i)
+
+      concrete_fn = worker_fn.get_concrete_function(x1)
+      del concrete_fn
+      del worker_fn
+
+    # No exception should be thrown from the thread
+    self._coord.join([t])
+
+  def testPendingNodesServerReplaced(self):
+    """Update cluster when nodes are still pending on remote workers."""
+    with ops.device(self.device_local):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    # Add enough ops so they are pending when changing the cluster
+    num_nodes = 10
+    ret = [None] * num_nodes
+    for i in range(num_nodes):
+      with ops.device(self.device_t1):
+        ret[i] = worker_fn(x1)
+    # While nodes are still pending on worker s1, replace worker s2 with s3.
+    context.update_server_def(server_def=self.server_def_s1_s3)
+    with ops.device(self.device_t2):
+      y = worker_fn(x1)
+    for i in range(num_nodes):
+      np.testing.assert_array_equal([[2, 2], [2, 2]], ret[i].numpy())
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testMultiThreadPendingNodesServerReplaced(self):
+    """Update cluster when other remote function calls are being launched."""
+    with ops.device(self.device_local):
+      x1 = array_ops.ones([2, 2])
+
+    num_calls = 10
+    lock = threading.Lock()
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    def thread_fn(device, results):
+      for i in range(num_calls):
+        lock.acquire()
+        with ops.device(device):
+          y = worker_fn(x1)
+        results[i] = y.numpy()
+        lock.release()
+
+    def update_server_def_fn():
+      for i in range(num_calls):
+        lock.acquire()
+        context.update_server_def(
+            server_def=(self.server_def_s1_s2 if i %
+                        2 == 0 else self.server_def_s1_s3))
+        lock.release()
+
+    t1_results = [None] * num_calls
+    t2_results = [None] * num_calls
+    threads = []
+    threads.append(
+        threading.Thread(target=thread_fn, args=(self.device_t1, t1_results)))
+    threads.append(
+        threading.Thread(target=thread_fn, args=(self.device_t2, t2_results)))
+    threads.append(threading.Thread(target=update_server_def_fn))
+    for t in threads:
+      t.start()
+    for t in threads:
+      t.join()
+    for result in t1_results + t2_results:
+      np.testing.assert_array_equal([[2, 2], [2, 2]], result)
+
+  def testMultiThreadPendingNodesLockFree(self):
+    """Update cluster when other remote function calls are being launched."""
+
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    num_calls = 10
+    self._coord = coordinator.Coordinator()
+
+    @def_function.function
+    def worker_fn(i):
+      return math_ops.matmul(i, i)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    def thread_fn(device, results):
+      for i in range(num_calls):
+        with self._coord.stop_on_exception():
+          with ops.device(device):
+            results[i] = worker_fn(x1).numpy()
+
+    def update_server_def_fn():
+      for _ in range(30):
+        with self._coord.stop_on_exception():
+          context.update_server_def(self.server_def_s1_s2)
+
+    t1_results = [None] * num_calls
+    t2_results = [None] * num_calls
+    threads = []
+    threads.append(
+        threading.Thread(target=thread_fn, args=(self.device_t1, t1_results)))
+    threads.append(
+        threading.Thread(target=thread_fn, args=(self.device_t2, t2_results)))
+    threads.append(threading.Thread(target=update_server_def_fn))
+    for t in threads:
+      t.start()
+    self._coord.join(threads)
+    for result in t1_results + t2_results:
+      np.testing.assert_array_equal([[2, 2], [2, 2]], result)
+
+  @test_util.run_in_async_and_sync_mode
+  def testDistributedFunctionServerAdded(self):
+    """Add a server to cluster, and run distributed function on it."""
+    with ops.device(self.device_t1):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      with ops.device(self.device_t2):
+        mul = math_ops.matmul(i, i)
+      return mul - array_ops.zeros_like(mul)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    context.update_server_def(server_def=self.server_def_s1_s2_s3)
+    with ops.device(self.device_t3):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testDistributedFunctionServerRemovedAddedBack(self):
+    """Add then remove a server, and run distributed function on cluster."""
+    with ops.device(self.device_local):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      with ops.device(self.device_t1):
+        mul = math_ops.matmul(i, i)
+      return mul - array_ops.zeros_like(mul)
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    context.update_server_def(server_def=self.server_def_s1)
+    with ops.device(self.device_t1):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    context.update_server_def(server_def=self.server_def_s1_s2)
+    with ops.device(self.device_t2):
+      y = worker_fn(x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testDistributedFunctionBothServersReplaced(self):
+    """Tests that replacing servers works correctly.
+
+    We create two servers, t1 and t2. We first replace t2, then we replace t1.
+
+    Among other things, this ensures that both already existing, and
+    restarted workers have the context view IDs correctly updated.
+    """
+    with ops.device(self.device_local):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      with ops.device(self.device_t1):
+        mul = math_ops.matmul(i, i)
+      with ops.device(self.device_t2):
+        add = mul + i
+      return add - i
+
+    # Forces function tracing and registration
+    worker_fn.get_concrete_function(x1)
+
+    # Replace task2
+    context.update_server_def(server_def=self.server_def_s1_s3)
+    for device in (self.device_t1, self.device_t2):
+      with ops.device(device):
+        y = worker_fn(x1)
+      np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # Then replace task1
+    context.update_server_def(server_def=self.server_def_s4_s3)
+    for device in (self.device_t1, self.device_t2):
+      with ops.device(device):
+        y = worker_fn(x1)
+      np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  def testDistributedFunctionPendingNodesServerReplaced(self):
+    with ops.device(self.device_local):
+      x1 = array_ops.ones([2, 2])
+
+    @def_function.function
+    def worker_fn(i):
+      with ops.device(self.device_t1):
+        mul = math_ops.matmul(i, i)
+      with ops.device(self.device_t2):
+        add = mul + i
+      return add - i
+
+    worker_fn.get_concrete_function(x1)
+
+    num_calls = 10
+    self._coord = coordinator.Coordinator()
+
+    def thread_fn(device, results):
+      with self._coord.stop_on_exception():
+        for i in range(num_calls):
+          with ops.device(device):
+            y = worker_fn(x1)
+          results[i] = y.numpy()
+
+    def update_server_def_fn():
+      with self._coord.stop_on_exception():
+        for i in range(num_calls):
+          context.update_server_def(
+              server_def=(self.server_def_s1_s2_s3 if i %
+                          2 == 0 else self.server_def_s1_s2))
+
+    results = [None] * num_calls
+    threads = []
+    threads.append(
+        threading.Thread(target=thread_fn, args=(self.device_t1, results)))
+    threads.append(threading.Thread(target=update_server_def_fn))
+    for t in threads:
+      t.start()
+    self._coord.join(threads)
+    for result in results:
+      np.testing.assert_array_equal([[2, 2], [2, 2]], result)
+
+  def testParameterServerMultiExecutors(self):
+    context.update_server_def(server_def=self.server_def_s1_s2_s3_s4)
+
+    with ops.device(self.device_t1):
+      v1 = variables.Variable(initial_value=0.)
+    with ops.device(self.device_t2):
+      v2 = variables.Variable(initial_value=10.)
+
+    @def_function.function
+    def worker_fn():
+      x1 = v1.read_value()
+      x2 = v2.read_value()
+      grad = (x1 + x2) * 0.1
+      v1.assign_add(grad)
+      v2.assign_sub(grad)
+      return v1 + v2
+
+    worker_fn.get_concrete_function()
+
+    executor_t3 = executor.new_executor(enable_async=False)
+    executor_t4 = executor.new_executor(enable_async=False)
+
+    num_calls = 10
+    self._coord = coordinator.Coordinator()
+
+    def thread_fn(executor_obj, device, results):
+      with self._coord.stop_on_exception():
+        for i in range(num_calls):
+          with context.executor_scope(executor_obj):
+            with ops.device(device):
+              results[i] = worker_fn()
+
+    def update_server_def_fn():
+      with self._coord.stop_on_exception():
+        for _ in range(30):
+          context.update_server_def(self.server_def_s1_s2_s3_s4)
+
+    t3_results = [None] * num_calls
+    t4_results = [None] * num_calls
+    threads = []
+    threads.append(
+        threading.Thread(
+            target=thread_fn, args=(executor_t3, self.device_t3, t3_results)))
+    threads.append(
+        threading.Thread(
+            target=thread_fn, args=(executor_t4, self.device_t4, t4_results)))
+    threads.append(threading.Thread(target=update_server_def_fn))
+    for t in threads:
+      t.start()
+    self._coord.join(threads)
+
+    # Cannot assert individual values since the results are non-deterministic.
+    # By summing up the value we ensure that there are all reasonable and valid
+    # numbers (not `None` or `NaN`).
+    total = np.sum(t3_results + t4_results)
+    self.assertGreater(total, 0)
+
+  def testCheckAlive(self):
+    with self.assertRaisesRegex(ValueError, "Context is not initialized."):
+      context.check_alive("/job:remote_device/task:0")
+    context.context().ensure_initialized()
+
+    self.assertTrue(context.check_alive("/job:remote_device/replica:0/task:0"))
+    self.assertTrue(context.check_alive("/job:remote_device/replica:0/task:1"))
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Unable to find worker interface"):
+      context.check_alive("/job:remote_device/replica:0/task:10")
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/remote_execution_test.py

@@ -0,0 +1,234 @@
+# Copyright 2018 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 remote eager execution."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+JOB_NAME = "remote_device"
+ALT_JOB_NAME = "alt_remote_device"
+
+
+def get_server_def(job_name, local_server_port, remote_server_addresses,
+                   task_index):
+  """Returns a server def with a single job + multiple tasks."""
+  cluster_def = cluster_pb2.ClusterDef()
+  job_def = cluster_def.job.add()
+  job_def.name = job_name
+  job_def.tasks[0] = "localhost:%d" % local_server_port
+
+  for i, remote_server_address in enumerate(remote_server_addresses, start=1):
+    job_def.tasks[i] = remote_server_address
+
+  server_def = tensorflow_server_pb2.ServerDef(
+      cluster=cluster_def,
+      job_name=job_name,
+      task_index=task_index,
+      protocol="grpc")
+
+  return server_def
+
+
+class RemoteExecutionTest(test.TestCase, parameterized.TestCase):
+
+  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
+    super(RemoteExecutionTest, self).__init__(methodName)
+    self._cached_server1 = server_lib.Server.create_local_server()
+    self._cached_server2 = server_lib.Server.create_local_server()
+
+    self._cached_server1_target = self._cached_server1.target[len("grpc://"):]
+    self._cached_server2_target = self._cached_server2.target[len("grpc://"):]
+
+  def setUp(self):
+    super(RemoteExecutionTest, self).setUp()
+    local_port = pywrap_tfe.TF_PickUnusedPortOrDie()
+    context.set_server_def(
+        server_def=get_server_def(
+            JOB_NAME,
+            local_server_port=local_port,
+            remote_server_addresses=[
+                self._cached_server1_target, self._cached_server2_target
+            ],
+            task_index=0))
+
+  def tearDown(self):
+    super(RemoteExecutionTest, self).tearDown()
+
+    # Clear the current device scope and reset the context to avoid polluting
+    # other test cases.
+    ops.device(None).__enter__()
+    context._reset_context()
+
+  @test_util.run_in_async_and_sync_mode
+  @test_util.run_gpu_only
+  def testGpuToRemoteCopy(self):
+    """Tests that the remote copy happens satisfactorily."""
+    x1 = array_ops.ones([2, 2]).gpu()
+    with ops.device("/job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      x2 = x1._copy()  # pylint: disable=protected-access
+
+    np.testing.assert_array_equal(x1.numpy(), x2.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  @test_util.run_gpu_only
+  def testGpuToRemoteOp(self):
+    with ops.device("gpu:0"):
+      x = array_ops.ones([2, 2])
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      y = math_ops.matmul(x, x)
+
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testDefunMatmul(self):
+    """Basic remote eager execution with defun."""
+
+    mm_defun = def_function.function(math_ops.matmul)
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      x1 = array_ops.ones([2, 2])
+    with ops.device("job:%s/replica:0/task:2/device:CPU:0" % JOB_NAME):
+      x2 = array_ops.ones([2, 2])
+      y = mm_defun(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testSimpleMatmul(self):
+    """Basic remote eager execution."""
+
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      x1 = array_ops.ones([2, 2])
+    with ops.device("job:%s/replica:0/task:2/device:CPU:0" % JOB_NAME):
+      x2 = array_ops.ones([2, 2])
+      y = math_ops.matmul(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  def testEagerPyFuncPlacement(self):
+    if not ops.executing_eagerly_outside_functions():
+      return
+
+    def f(x):
+      return math_ops.square(x)
+
+    with ops.device("/job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      const_op = constant_op.constant(3.0, dtype=dtypes.float32)
+      # PyFuncOp should be placed on the localhost's address space.
+      py_func_op = script_ops.eager_py_func(
+          func=f, inp=[const_op], Tout=dtypes.float32)
+      self.assertEqual(py_func_op.device,
+                       "/job:%s/replica:0/task:0/device:CPU:0" % JOB_NAME)
+      self.assertEqual(self.evaluate(py_func_op), 9.0)
+
+  @test_util.run_in_async_and_sync_mode
+  def testSimpleWeightRead(self):
+    """Basic remote eager weight read."""
+
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      w = resource_variable_ops.ResourceVariable([[2.0]])
+      loss = w * w
+    np.testing.assert_array_equal([[4.0]], loss.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testTapeWeightRead(self):
+    """Remote eager weight read in a tape."""
+
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      w = resource_variable_ops.ResourceVariable([[3.0]])
+      with backprop.GradientTape() as tape:
+        loss = w * w
+
+      grad = tape.gradient(loss, w)
+    np.testing.assert_array_equal([[9.0]], loss.numpy())
+    np.testing.assert_array_equal([[6.0]], grad.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testServerDefChanged(self):
+    """Update server def, and run ops on new cluster."""
+    context.set_server_def(
+        server_def=get_server_def(
+            ALT_JOB_NAME,
+            local_server_port=0,
+            remote_server_addresses=[
+                self._cached_server1_target, self._cached_server2_target
+            ],
+            task_index=0))
+
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % ALT_JOB_NAME):
+      x1 = array_ops.ones([2, 2])
+    y = math_ops.matmul(x1, x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # Set the server def back to JOB_NAME
+    context.set_server_def(
+        server_def=get_server_def(
+            JOB_NAME,
+            local_server_port=0,
+            remote_server_addresses=[
+                self._cached_server1_target, self._cached_server2_target
+            ],
+            task_index=0))
+
+    with ops.device("job:%s/replica:0/task:1/device:CPU:0" % JOB_NAME):
+      x1 = array_ops.ones([2, 2])
+    y = math_ops.matmul(x1, x1)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testConnectToRemoteServer(self):
+    """Basic server connection."""
+    context._reset_context()
+    remote.connect_to_remote_host(self._cached_server1_target)
+
+    with ops.device("job:worker/replica:0/task:0/device:CPU:0"):
+      x1 = array_ops.ones([2, 2])
+      x2 = array_ops.ones([2, 2])
+      y = math_ops.matmul(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+  @test_util.run_in_async_and_sync_mode
+  def testContextDeviceUpdated(self):
+    """Tests that the context device is correctly updated."""
+
+    with ops.device("cpu:0"):
+      x1 = array_ops.ones([2, 2])
+      x2 = array_ops.ones([2, 2])
+      y = math_ops.matmul(x1, x2)
+    np.testing.assert_array_equal([[2, 2], [2, 2]], y.numpy())
+
+    # `y` is placed on the local CPU as expected.
+    self.assertEqual(y.device,
+                     "/job:%s/replica:0/task:0/device:CPU:0" % JOB_NAME)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/eager/remote_test.py

@@ -0,0 +1,808 @@
+# Copyright 2019 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 remote execution."""
+
+import os
+import random
+import time
+
+from absl.testing import parameterized
+import numpy as np
+import portpicker
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver
+from tensorflow.python.eager import cancellation
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import executor
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.training import server_lib
+from tensorflow.python.training.server_lib import ClusterSpec
+from tensorflow.python.util import compat
+
+
+class SingleWorkerTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(SingleWorkerTest, self).setUp()
+
+    workers, _ = test_util.create_local_cluster(1, 0)
+    remote.connect_to_remote_host(workers[0].target)
+
+  def tearDown(self):
+    super(SingleWorkerTest, self).tearDown()
+
+    # Clear the current device scope to avoid polluting other test cases.
+    ops.device(None).__enter__()
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  def testMultiDeviceFunctionBasic(self):
+
+    @def_function.function
+    def basic(i):
+      with ops.device('/job:localhost/replica:0/task:0/cpu:0'):
+        a = constant_op.constant([2]) + i
+      with ops.device('/job:worker/replica:0/task:0/cpu:0'):
+        b = constant_op.constant([1])
+
+      return a + b
+
+    self.assertAllEqual(basic(constant_op.constant([2])).numpy(), [5])
+    self.assertAllEqual(basic(constant_op.constant([1])).numpy(), [4])
+
+  def testMultiDeviceFunctionVariable(self):
+    with ops.device('/job:worker/replica:0/task:0/cpu:0'):
+      variable_b = variables.Variable(1)
+
+    # Add a sync point to avoid the out-of-order issue of eager async execution
+    # (b/155789951).
+    context.async_wait()
+
+    @def_function.function
+    def with_variable(i):
+      return i + variable_b
+
+    self.assertAllEqual(with_variable(constant_op.constant([2])).numpy(), [3])
+
+  def testMultiDeviceFunctionRemoteOutput(self):
+    with ops.device('/job:worker/replica:0/task:0/cpu:0'):
+      variable_b = variables.Variable(1)
+
+    @def_function.function
+    def remote_output(i):
+      with ops.device('/job:worker/replica:0/task:0/cpu:0'):
+        c = variable_b + 1
+      return i + variable_b, c
+
+    rets = remote_output(constant_op.constant([1]))
+    self.assertAllEqual(rets[0].numpy(), [2])
+    self.assertAllEqual(rets[1].numpy(), 2)
+    self.assertEqual(rets[0].backing_device,
+                     '/job:localhost/replica:0/task:0/device:CPU:0')
+    self.assertEqual(rets[1].backing_device,
+                     '/job:worker/replica:0/task:0/device:CPU:0')
+
+  def testStreaming(self):
+    """A mini stress test for streaming - issuing many RPCs back to back."""
+    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
+      x = array_ops.ones([2, 2])
+      y = array_ops.zeros([2, 2])
+      num_iters = 200
+      for _ in range(num_iters):
+        y = x + y
+        # Ask for y's shape after every 10 additions on average.
+        # This exercises waiting for remote shape logic in TensorHandle.
+        if random.randint(1, 10) == 1:
+          _ = y.shape
+    np.testing.assert_array_equal(
+        [[num_iters, num_iters], [num_iters, num_iters]], y.numpy())
+
+  def testTwoExecutors(self):
+    # Run an op on the main executor that by default uses StreamingEnqueue to
+    # schedule the op to run on the remote async executor. This op produces an
+    # error, i.e., division by zero, but will not be immediately caught due to
+    # streaming enqueue.
+    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
+      a = constant_op.constant(3)
+      b = constant_op.constant(0)
+      math_ops.div(a, b)
+
+    # Run another op using another executor that disables streaming enqueue,
+    # which would run the op using the tf_compute thread pool in the remote
+    # worker. Since the op is not run in the same remotes async executor, it
+    # will not carry back that error produced by the op above, even though this
+    # op is executed synchronously.
+    with context.executor_scope(
+        executor.new_executor(
+            enable_async=False, enable_streaming_enqueue=False)):
+      with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
+        c = constant_op.constant(4)
+        d = constant_op.constant(2)
+        self.assertEqual(math_ops.div(c, d).numpy(), 2)
+
+    # Sync on the context to force to catch the error produced by the first op.
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      context.async_wait()
+    self.assertIn('division by zero', cm.exception.message)
+
+  def testShapeError_OpByOp(self):
+    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
+      x = array_ops.ones([2, 3])
+      y = array_ops.zeros([2, 2])
+      with self.assertRaises(errors.InvalidArgumentError) as cm:
+        math_ops.matmul(x, y)
+
+    self.assertIn('Dimensions must be equal', cm.exception.message)
+
+  def testShapeError_Function(self):
+
+    @def_function.function
+    def matmul_func(x, y):
+      return math_ops.matmul(x, y)
+
+    x = array_ops.ones([2, 3])
+    y = array_ops.zeros([2, 2])
+
+    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
+      with self.assertRaises(ValueError) as cm:
+        matmul_func(x, y)
+
+    self.assertIn('Dimensions must be equal', cm.exception.args[0])
+
+  def testClientVarible(self):
+    var = variables.Variable(initial_value=0)
+
+    @def_function.function
+    def func():
+      with ops.device('/job:localhost/task:0'):
+        read = var.read_value()
+      return read + 1
+
+    with ops.device('/job:worker/task:0'):
+      self.assertAllEqual(func(), 1)
+
+  def testRemoteCall(self):
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def _remote_fn(x):
+      return constant_op.constant(1) + x
+
+    remote_fn = _remote_fn.get_concrete_function()
+
+    @def_function.function
+    def func(x):
+      return functional_ops.remote_call(
+          args=[x],
+          Tout=[dtypes.int32],
+          f=remote_fn,
+          target='/job:worker/task:0')
+
+    with ops.device('/job:localhost/task:0'):
+      self.assertAllEqual(func(constant_op.constant(1)), [2])
+
+  def testOperationTimeout(self):
+    context._reset_context()
+    context.context().operation_timeout_in_ms = 10
+    workers, _ = test_util.create_local_cluster(1, 0)
+    remote.connect_to_remote_host(workers[0].target)
+
+    q = data_flow_ops.FIFOQueue(1, dtypes.int32)
+
+    @def_function.function
+    def f():
+      return q.dequeue()
+
+    with self.assertRaises(errors.DeadlineExceededError):
+      with ops.device('/job:worker/replica:0/task:0'):
+        f()
+      # If streaming RPC is enabled, fetch remote errors before end of execution
+      context.async_wait()
+
+
+class RemoteAsyncTest(test.TestCase):
+
+  def setUp(self):
+    super(RemoteAsyncTest, self).setUp()
+
+    workers, _ = test_util.create_local_cluster(1, 0)
+    remote.connect_to_remote_host(workers[0].target)
+
+  def tearDown(self):
+    super(RemoteAsyncTest, self).tearDown()
+
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  def test_out_of_range_with_while_loop(self):
+
+    with ops.device('/job:worker/task:0'):
+      dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0])
+      dataset = dataset.batch(1, drop_remainder=False)
+      iterator = iter(dataset)
+      v = variables.Variable(1.0)
+
+    @def_function.function
+    def train_step(iterator):
+      i = next(iterator)
+      v.assign_add(math_ops.reduce_mean(i))
+
+    while True:
+      try:
+        with ops.device('/job:worker/task:0'):
+          train_step(iterator)
+      except (errors.OutOfRangeError, errors.InternalError):
+        context.async_clear_error()
+        break
+
+    self.assertAllEqual(v.numpy(), 4.0)
+
+  def test_out_of_range_with_for_loop(self):
+
+    with ops.device('/job:worker/task:0'):
+      dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0])
+      dataset = dataset.batch(1, drop_remainder=False)
+      iterator = iter(dataset)
+      v = variables.Variable(1.0)
+
+    @def_function.function
+    def train_step(iterator):
+      i = next(iterator)
+      v.assign_add(math_ops.reduce_mean(i))
+
+    num_steps = 3
+    for i in range(num_steps):
+      try:
+        with ops.device('/job:worker/task:0'):
+          train_step(iterator)
+        if i == num_steps - 1:
+          context.async_wait()
+      except errors.OutOfRangeError:
+        context.async_clear_error()
+        break
+
+    self.assertAllEqual(v.numpy(), 4.0)
+
+  def test_out_of_range_with_async_scope(self):
+
+    with ops.device('/job:worker/task:0'):
+      dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0])
+      dataset = dataset.batch(1, drop_remainder=False)
+      iterator = iter(dataset)
+      v = variables.Variable(1.0)
+
+    @def_function.function
+    def train_step(iterator):
+      i = next(iterator)
+      v.assign_add(math_ops.reduce_mean(i))
+
+    num_steps = 3
+    try:
+      with context.async_scope():
+        for _ in range(num_steps):
+          with ops.device('/job:worker/task:0'):
+            train_step(iterator)
+    except errors.OutOfRangeError:
+      context.async_clear_error()
+
+    self.assertAllEqual(v.numpy(), 4.0)
+
+
+class MultiWorkersTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(MultiWorkersTest, self).setUp()
+
+    workers, _ = test_util.create_local_cluster(3, 0)
+    remote.connect_to_remote_host(
+        [workers[0].target, workers[1].target, workers[2].target])
+
+  def tearDown(self):
+    super(MultiWorkersTest, self).tearDown()
+
+    # Clear the current device scope to avoid polluting other test cases.
+    ops.device(None).__enter__()
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  def testReturnRemoteArgument(self):
+
+    @def_function.function
+    def local_func(i):
+      return i
+
+    with ops.device('/job:worker/replica:0/task:0'):
+      x = constant_op.constant([2, 1])
+
+    with ops.device('/job:worker/replica:0/task:1'):
+      self.assertAllEqual(local_func(x), [2, 1])
+
+  def testMultiDeviceFunctionAmbiguousDevice(self):
+
+    @def_function.function
+    def ambiguous_device(i):
+      with ops.device('/job:worker'):
+        # Multiple worker tasks, thus ambiguous device found error will be
+        # raised.
+        return i + constant_op.constant([2])
+
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      ambiguous_device(constant_op.constant([2])).numpy()
+
+    self.assertIn('the output node must match exactly one device',
+                  cm.exception.message)
+
+  # Note that the following tests for remote function cancellation only works
+  # when non-streaming RPC. We need to disable streaming explicitly and restore
+  # this config to its initial value at the end of each test case.
+  def testCancelRemoteFunctionBeforeExecution(self):
+    remote_async_env_var = 'TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE'
+    default_streaming = os.environ.get(remote_async_env_var)
+    os.environ[remote_async_env_var] = str(False)
+
+    q = data_flow_ops.FIFOQueue(1, dtypes.int32)
+
+    @def_function.function
+    def f():
+      return q.dequeue()
+
+    c_mgr = cancellation.CancellationManager()
+    cancelable_func = c_mgr.get_cancelable_function(f.get_concrete_function())
+
+    c_mgr.start_cancel()
+    with self.assertRaises(errors.CancelledError):
+      with ops.device('/job:worker/replica:0/task:1'):
+        cancelable_func()
+
+    if default_streaming is None:
+      del os.environ[remote_async_env_var]
+    else:
+      os.environ[remote_async_env_var] = default_streaming
+
+  def testCancelRemoteFunctionDuringExecution(self):
+    remote_async_env_var = 'TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE'
+    default_streaming = os.environ.get(remote_async_env_var)
+    os.environ[remote_async_env_var] = str(False)
+
+    q = data_flow_ops.FIFOQueue(1, dtypes.int32)
+
+    @def_function.function
+    def f():
+      return q.dequeue()
+
+    c_mgr = cancellation.CancellationManager()
+    cancelable_func = c_mgr.get_cancelable_function(f.get_concrete_function())
+
+    def cancel_thread():
+      time.sleep(0.5)
+      c_mgr.start_cancel()
+
+    t = self.checkedThread(cancel_thread)
+    t.start()
+    with self.assertRaises(errors.CancelledError):
+      with ops.device('/job:worker/replica:0/task:1'):
+        cancelable_func()
+    t.join()
+
+    if default_streaming is None:
+      del os.environ[remote_async_env_var]
+    else:
+      os.environ[remote_async_env_var] = default_streaming
+
+  def testMultiDeviceFunctionOnLocalDevice(self):
+    with ops.device('/job:worker/replica:0/task:1'):
+      variable_b = variables.Variable(1.0)
+
+    @def_function.function
+    def remote_function(i):
+      with ops.device('/job:worker/replica:0/task:0'):
+        a = i + variable_b
+      c = a + 1.0
+      return c
+
+    self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])
+
+  def testMultiDeviceFunctionExecutionOrderingWithPackedInput(self):
+    shape = [2]
+    with ops.device('/job:worker/replica:0/task:2/device:CPU:0'):
+      # Send 20 remote requests to simulate heavy load on worker:2.
+      unused_values = []
+      for _ in range(20):
+        unused_values.append(array_ops.zeros(shape))
+      func_input = array_ops.zeros(shape)
+
+    packed_input = ops.pack_eager_tensors([func_input])
+
+    @def_function.function
+    def func(packed_input):
+      # When worker:2 receives the component function request, packed_input
+      # should be ready on worker:2.
+      with ops.device('/job:worker/replica:0/task:2/device:CPU:0'):
+        ret = packed_input + constant_op.constant(1.0)
+      return ret + constant_op.constant(1.0)
+
+    # Run the function on a worker:1
+    with ops.device('/job:worker/replica:0/task:1/device:CPU:0'):
+      self.assertAllEqual(func(packed_input).numpy(),
+                          array_ops.ones(shape).numpy() * 2)
+
+  def testMultiDeviceFunctionWithPackedVariable(self):
+    with ops.device('/job:worker/replica:0/task:0/device:CPU:0'):
+      var0 = resource_variable_ops.ResourceVariable(1.0)
+    with ops.device('/job:worker/replica:0/task:1/device:CPU:0'):
+      var1 = resource_variable_ops.ResourceVariable(2.0)
+
+    packed_var = ops.pack_eager_tensors([var0.handle, var1.handle])
+    self.assertEqual(packed_var.device,
+                     '/job:localhost/replica:0/task:0/device:COMPOSITE:0')
+    self.assertEqual(packed_var.backing_device,
+                     '/job:localhost/replica:0/task:0/device:COMPOSITE:0')
+
+    @def_function.function
+    def add_variables():
+      with ops.device('/job:worker/replica:0/task:0/device:CPU:0'):
+        read0 = resource_variable_ops.read_variable_op(
+            packed_var, dtype=dtypes.float32)
+      with ops.device('/job:worker/replica:0/task:1/device:CPU:0'):
+        read1 = resource_variable_ops.read_variable_op(
+            packed_var, dtype=dtypes.float32)
+
+      return read0 + read1
+
+    # Run the function on a remote device
+    with ops.device('/job:worker/replica:0/task:0'):
+      self.assertAllEqual(add_variables().numpy(), 3.0)
+
+    # Run the function on a local worker
+    self.assertAllEqual(add_variables().numpy(), 3.0)
+
+  def testMultiDeviceFunctionOnRemoteDeviceWithWait(self):
+    with ops.device('/job:worker/replica:0/task:1'):
+      variable_b = variables.Variable([1.0])
+
+    @def_function.function
+    def remote_function(i):
+      x = array_ops.ones([1000, 1000])
+      for _ in range(1, 1000):
+        x = x * x
+      variable_b.assign_add(i)
+      a = 1.0 + variable_b
+      return a
+
+    @def_function.function
+    def remote_function2(i):
+      variable_b.assign_add(i)
+      a = 1.0 + variable_b
+      return a
+
+    # Runs first function:
+    # - on remote device
+    # - needs remote input
+    # - is side impacting
+    # - runs much slower
+    with ops.device('/job:worker/replica:0/task:0'):
+      remote_function(constant_op.constant([2.0]))
+
+    # Runs second function:
+    # - on remote device
+    # - is side impacting
+    # There should be a sync point here and the next function will be executed
+    # only after the first function has completed.
+    with ops.device('/job:worker/replica:0/task:2'):
+      self.assertAllEqual(remote_function2(constant_op.constant([3.0])), [7.0])
+
+  def testMultiDeviceFunctionOnRemoteDevice(self):
+    with ops.device('/job:worker/replica:0/task:1'):
+      variable_b = variables.Variable(1.0)
+
+    @def_function.function
+    def remote_function(i):
+      with ops.device('/job:worker/replica:0/task:0'):
+        a = i + variable_b
+      c = a + 1.0
+      return c
+
+    with ops.device('/job:worker/replica:0/task:0'):
+      self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])
+
+    if test_util.is_gpu_available():
+      with ops.device('/job:worker/replica:0/task:0/device:GPU:0'):
+        self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])
+
+  def testMultiDeviceFunctionRemoteOutput(self):
+    with ops.device('/job:worker/replica:0/task:1/cpu:0'):
+      variable_b = variables.Variable(1)
+
+    @def_function.function
+    def remote_output(i):
+      with ops.device('/job:worker/replica:0/task:1/cpu:0'):
+        c = variable_b + 1
+      return i + variable_b, c
+
+    with ops.device('/job:worker/replica:0/task:0/cpu:0'):
+      rets = remote_output(constant_op.constant([1]))
+    self.assertEqual(rets[0].backing_device,
+                     '/job:worker/replica:0/task:0/device:CPU:0')
+    self.assertEqual(rets[1].backing_device,
+                     '/job:worker/replica:0/task:1/device:CPU:0')
+    self.assertAllEqual(rets[0].numpy(), [2])
+    self.assertAllEqual(rets[1].numpy(), 2)
+
+  def testMultiDeviceWhileLoopOnRemoteDevice(self):
+    with ops.device('/job:worker/replica:0/task:1'):
+      variable_b = variables.Variable(1.0)
+
+    @def_function.function
+    def remote_function(i):
+
+      def body(i, _):
+        with ops.device('/job:worker/replica:0/task:0'):
+          a = i + variable_b
+        return a + 1.0, 1
+
+      return while_loop.while_loop_v2(lambda _, d: d < 1, body, [i, 0])[0]
+
+    with ops.device('/job:worker/replica:0/task:0'):
+      self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])
+
+    if test_util.is_gpu_available():
+      with ops.device('/job:worker/replica:0/task:0/device:GPU:0'):
+        self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])
+
+  def testSimpleParameterServer(self):
+
+    with ops.device('/job:worker/task:2/device:CPU:0'):
+      v1 = variables.Variable(initial_value=0)
+      v2 = variables.Variable(initial_value=10)
+
+    @def_function.function
+    def worker_fn():
+      v1.assign_add(1)
+      v2.assign_sub(2)
+      return v1.read_value() + v2.read_value()
+
+    with ops.device('/job:worker/task:0/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 9)
+
+    with ops.device('/job:worker/task:1/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 8)
+
+
+_GRPC_PREFIX = 'grpc://'
+
+
+class MultiJobsTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(MultiJobsTest, self).setUp()
+
+    workers, ps = test_util.create_local_cluster(num_workers=2, num_ps=2)
+    cluster = {
+        'my_worker': [_strip_prefix(t.target, _GRPC_PREFIX) for t in workers],
+        'my_ps': [_strip_prefix(t.target, _GRPC_PREFIX) for t in ps],
+    }
+    self._cluster = server_lib.ClusterSpec(cluster)
+    self._cluster_resolver = SimpleClusterResolver(
+        cluster_spec=self._cluster, master=ps[0].target)
+
+  def tearDown(self):
+    super(MultiJobsTest, self).tearDown()
+
+    # Clear the current device scope to avoid polluting other test cases.
+    ops.device(None).__enter__()
+    # Reset the context to avoid polluting other test cases.
+    context._reset_context()
+
+  def testMultipleDeviceFoundCheck(self):
+    remote.connect_to_cluster(self._cluster)
+
+    @def_function.function
+    def func():
+      with ops.device('cpu:0'):
+        # Multiple CPU:0 devices match would be found, but the CPU:0 from the
+        # parent device scope should be picked.
+        x = test_ops.device_placement_op()
+        y = string_ops.string_upper(x)
+        packed_var_0 = array_ops_stack.stack([x, y], 0)
+        return packed_var_0
+
+    with ops.device('/job:my_worker/task:1'):
+      output = self.evaluate(func())
+      self.assertEqual(
+          compat.as_bytes('/job:my_worker/replica:0/task:1/device:CPU:0'),
+          output[0])
+      self.assertIn(compat.as_bytes('/JOB:MY_WORKER'), output[1])
+    with ops.device('/job:my_ps/task:1'):
+      output = self.evaluate(func())
+      self.assertEqual(
+          compat.as_bytes('/job:my_ps/replica:0/task:1/device:CPU:0'),
+          output[0])
+      self.assertIn(compat.as_bytes('/JOB:MY_PS'), output[1])
+
+  def testSimpleParameterServer(self):
+    remote.connect_to_cluster(self._cluster)
+
+    with ops.device('/job:my_ps/task:0/device:CPU:0'):
+      v1 = variables.Variable(initial_value=0)
+      v2 = variables.Variable(initial_value=10)
+
+    @def_function.function
+    def worker_fn():
+      v1.assign_add(1)
+      v2.assign_sub(2)
+      return v1.read_value() + v2.read_value()
+
+    with ops.device('/job:my_worker/task:0/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 9)
+
+    with ops.device('/job:my_worker/task:1/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 8)
+
+  def testResetClusterWithDifferentJobNames(self):
+    addr = 'localhost:%s' % portpicker.pick_unused_port()
+    cluster = server_lib.ClusterSpec({'localhost': [addr]})
+    remote.connect_to_cluster(cluster, job_name='localhost')
+    with ops.device('/job:localhost/task:0/device:CPU:0'):
+      v1 = variables.Variable(initial_value=0)
+      v1.assign_add(1)
+
+    # Replace job name from 'localhost' to 'worker' in the cluster.
+    addr = 'localhost:%s' % portpicker.pick_unused_port()
+    cluster = server_lib.ClusterSpec({'worker': [addr]})
+    remote.connect_to_cluster(cluster, job_name='worker')
+
+    with ops.device('/job:worker/task:0/device:CPU:0'):
+      v2 = variables.Variable(initial_value=0)
+      v2.assign_add(1)
+
+  # TODO(b/152224115): Re-enable this test.
+  def DISABLED_testSimpleParameterServerWithDeviceFilters(self):
+    cluster_device_filters = server_lib.ClusterDeviceFilters()
+    for i in range(2):
+      cluster_device_filters.set_device_filters('my_worker', i, ['/job:my_ps'])
+      cluster_device_filters.set_device_filters('my_ps', i, ['/job:my_worker'])
+    remote.connect_to_cluster(
+        self._cluster, cluster_device_filters=cluster_device_filters)
+
+    with ops.device('/job:my_ps/task:0/device:CPU:0'):
+      v1 = variables.Variable(initial_value=0)
+    with ops.device('/job:my_ps/task:1/device:CPU:0'):
+      v2 = variables.Variable(initial_value=10)
+
+    @def_function.function
+    def worker_fn():
+      v1.assign_add(1)
+      v2.assign_sub(2)
+      return v1.read_value() + v2.read_value()
+
+    with ops.device('/job:my_worker/task:0/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 9)
+    with ops.device('/job:my_worker/task:1/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 8)
+
+    # The following remote call would fail because the ps nodes cannot see each
+    # other due to the device filters.
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      with ops.device('/job:my_ps/task:0/device:CPU:0'):
+        worker_fn().numpy()
+    self.assertIn('/job:my_ps/replica:0/task:1/device:CPU:0 unknown device',
+                  cm.exception.message)
+
+    with self.assertRaises(errors.InvalidArgumentError) as cm:
+      with ops.device('/job:my_ps/task:1/device:CPU:0'):
+        worker_fn().numpy()
+    self.assertIn('/job:my_ps/replica:0/task:0/device:CPU:0 unknown device',
+                  cm.exception.message)
+
+    with ops.device('/job:my_worker/task:0/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 7)
+    with ops.device('/job:my_worker/task:1/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 6)
+    # Explicitly delete variables to avoid triggering errors when being GC'ed in
+    # subsequent tests.
+    del v1, v2
+
+  def testConnectWithClusterResolver(self):
+    remote.connect_to_cluster(self._cluster_resolver)
+
+    v1 = variables.Variable(initial_value=0)
+    v2 = variables.Variable(initial_value=10)
+
+    @def_function.function
+    def worker_fn():
+      v1.assign_add(1)
+      v2.assign_sub(2)
+      return v1.read_value() + v2.read_value()
+
+    with ops.device('/job:my_worker/task:0/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 9)
+
+    with ops.device('/job:my_worker/task:1/device:CPU:0'):
+      self.assertAllEqual(worker_fn(), 8)
+
+  def testConnectToClusterTwiceOk(self):
+    remote.connect_to_cluster(self._cluster_resolver)
+    remote.connect_to_cluster(self._cluster_resolver)
+
+  def testConnectToClusterOnMismatchedDevice(self):
+    remote.connect_to_cluster(self._cluster_resolver)
+
+    # enter into another device scope.
+    ops.device('/job:my_worker/task:0/device:CPU:0').__enter__()
+
+    with self.assertRaises(ValueError):
+      remote.connect_to_cluster(self._cluster_resolver)
+
+  def testConnectToClusterWithLocalMaster(self):
+    local_resolver = SimpleClusterResolver(ClusterSpec({}), master='local')
+    remote.connect_to_cluster(local_resolver)
+
+  def testConnectToClusterInGraphModeWillFail(self):
+    ops.disable_eager_execution()
+    with self.assertRaises(ValueError):
+      remote.connect_to_cluster(self._cluster_resolver)
+    ops.enable_eager_execution()
+
+  def testConnectToClusterWithoutLocalGpu(self):
+    # Only remote workers have GPU devices
+    context.context().set_visible_devices([], 'GPU')
+    # Ensure that no default device is set in eager context
+    remote.connect_to_cluster(self._cluster_resolver,
+                              make_master_device_default=False)
+    self.assertEmpty(context.get_device_name())
+
+    v1 = variables.Variable(initial_value=0)
+    v1.assign_add(1)
+    self.assertAllEqual(v1.read_value(), 1)
+
+  # TODO(b/249134783): Add a test for task failures by introducing an Op for
+  # reporting errors.
+  def testGetTaskStatesAllOK(self):
+    context.context().configure_coordination_service(
+        service_type='standalone', service_leader='/job:my_ps/replica:0/task:0')
+    remote.connect_to_cluster(self._cluster)
+    context.context().ensure_initialized()
+
+    states = context.context().get_task_states([('my_worker', 2), ('my_ps', 2)])
+    self.assertLen(states, 4)
+    for state in states:
+      self.assertIsNone(state)
+
+
+def _strip_prefix(s, prefix):
+  return s[len(prefix):] if s.startswith(prefix) else s
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/eager/run_eager_op_as_function_test.py

@@ -0,0 +1,277 @@
+# Copyright 2021 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 wrapping an eager op in a call op at runtime."""
+import time
+
+from tensorflow.python.data.experimental.ops import prefetching_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import benchmarks_test_base
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import bitwise_ops
+from tensorflow.python.ops import critical_section_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_map_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import tf_inspect
+
+
+def run_benchmark(func, num_iters, unused_execution_mode):
+  # warm up
+  func()
+  start = time.time()
+  for _ in range(num_iters):
+    func()
+  end = time.time()
+  return end - start
+
+
+CPU = "/device:CPU:0"
+GPU = "/device:GPU:0"
+
+
+# TODO(srbs): Why can't we use absl parameterized here?
+@test_util.with_eager_op_as_function
+class MicroBenchmarks(benchmarks_test_base.MicroBenchmarksBase):
+
+  def __init__(self):
+    super().__init__()
+    self._m_2_by_2 = random_ops.random_uniform((2, 2))
+    self._m_2_by_2_int32 = random_ops.random_uniform((2, 2),
+                                                     maxval=5,
+                                                     dtype=dtypes.int32)
+    self._m_100_by_100 = random_ops.random_uniform((100, 100))
+    self._m_100_by_100_int32 = random_ops.random_uniform((100, 100),
+                                                         maxval=5,
+                                                         dtype=dtypes.int32)
+    self._m_1000_by_1000 = random_ops.random_uniform((1000, 1000))
+    self._m_1000_by_1000_int32 = random_ops.random_uniform((1000, 1000),
+                                                           maxval=5,
+                                                           dtype=dtypes.int32)
+
+  def _get_benchmark_name(self):
+    """Copied from benchmarks_test.py."""
+    stack = tf_inspect.stack()
+    name = None
+    for frame in stack[::-1]:
+      f_locals = frame[0].f_locals
+      f_self = f_locals.get("self", None)
+      if isinstance(f_self, test.Benchmark):
+        name = frame[3]  # Get the method name
+        # This is a hack to get around the fact that some methods might have a
+        # disable_tfrt decorator around them. In that case a function called
+        # 'decorated' wraps the real called function underneath and so we
+        # peek one deeper into the stack to get the real name.
+        if name == "decorated":
+          continue
+        else:
+          break
+    if name is None:
+      raise ValueError("Unable to determine calling Benchmark function.")
+    if context.is_tfrt_enabled():
+      name = name + "_tfrt"
+    if context.run_eager_op_as_function_enabled():
+      name = name + "_eager_op_as_function"
+    return name
+
+  def _run(self, func, num_iters):
+    self.run_report(run_benchmark, func, num_iters, report_mean_us=True)
+
+  def _benchmark_matmul(self, mat, device):
+    if device == GPU and not context.num_gpus():
+      return
+    with context.device(device):
+      if device == GPU:
+        mat = mat.gpu()
+      func = lambda: math_ops.matmul(mat, mat)
+      self._run(func, num_iters=5000)
+
+  def _benchmark_bitwise_and(self, mat, device):
+    if device == GPU and not context.num_gpus():
+      return
+    with context.device(device):
+      if device == GPU:
+        mat = mat.gpu()
+      func = lambda: bitwise_ops.bitwise_and(mat, mat)
+      self._run(func, num_iters=5000)
+
+  def _benchmark_random_normal(self, device):
+    if device == GPU and not context.num_gpus():
+      return
+    with context.device(device):
+
+      def func():
+        mat = constant_op.constant([3], dtypes.int32)
+        s = mat + mat
+        random_ops.random_normal(shape=s)
+
+      self._run(func, num_iters=5000)
+
+  # This only needs to be tested on GPU where redundant data transfers can occur
+  def benchmark_random_normal_GPU(self):
+    self._benchmark_random_normal(GPU)
+
+  def benchmark_tf_matmul_2_by_2_CPU(self):
+    self._benchmark_matmul(self._m_2_by_2, CPU)
+
+  def benchmark_tf_bitwise_and_2_by_2_CPU(self):
+    self._benchmark_bitwise_and(self._m_2_by_2_int32, CPU)
+
+  def benchmark_tf_matmul_2_by_2_GPU(self):
+    self._benchmark_matmul(self._m_2_by_2, GPU)
+
+  def benchmark_tf_bitwise_and_2_by_2_GPU(self):
+    self._benchmark_bitwise_and(self._m_2_by_2_int32, GPU)
+
+  def benchmark_tf_matmul_100_by_100_CPU(self):
+    self._benchmark_matmul(self._m_100_by_100, CPU)
+
+  def benchmark_tf_bitwise_and_100_by_100_CPU(self):
+    self._benchmark_bitwise_and(self._m_100_by_100_int32, CPU)
+
+  def benchmark_tf_matmul_100_by_100_GPU(self):
+    self._benchmark_matmul(self._m_100_by_100, GPU)
+
+  def benchmark_tf_bitwise_and_100_by_100_GPU(self):
+    self._benchmark_bitwise_and(self._m_100_by_100_int32, GPU)
+
+  def benchmark_tf_matmul_1000_by_1000_CPU(self):
+    self._benchmark_matmul(self._m_1000_by_1000, CPU)
+
+  def benchmark_tf_bitwise_and_1000_by_1000_CPU(self):
+    self._benchmark_bitwise_and(self._m_1000_by_1000_int32, CPU)
+
+  def benchmark_tf_matmul_1000_by_1000_GPU(self):
+    self._benchmark_matmul(self._m_1000_by_1000, GPU)
+
+  def benchmark_tf_bitwise_and_1000_by_1000_GPU(self):
+    self._benchmark_bitwise_and(self._m_1000_by_1000_int32, GPU)
+
+
+@test_util.with_eager_op_as_function
+class RunEagerOpAsFunctionTest(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._m_2_by_2 = random_ops.random_uniform((2, 2))
+
+  def testDefaultAttrValues(self):
+    ragged_map_ops.map_fn(
+        fn=lambda x: x,
+        elems=ragged_factory_ops.constant([[7]]),
+        dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int32, ragged_rank=1))
+
+  def testArrayFill(self):
+    array_ops.fill(
+        constant_op.constant([2], dtype=dtypes.int64), constant_op.constant(1))
+
+  def testDatasetMap(self):
+    # When a GPU is available, this would test that the wrapped call ops are
+    # placed on the CPU (i.e. the device is selected using the unwrapped op).
+    dataset_ops.Dataset.range(2).map(math_ops.square)
+
+  def testPrefetchToDevice(self):
+    if not context.num_gpus():
+      self.skipTest("No GPU available")
+    dataset = dataset_ops.Dataset.range(10)
+    dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
+
+  def testMatmul(self):
+    math_ops.matmul(self._m_2_by_2, self._m_2_by_2)
+
+  def testMixedTypeListInputFastPath(self):
+    array_ops.identity_n([self._m_2_by_2, self._m_2_by_2])
+
+  def testMixedTypeListInputEagerFallback(self):
+    array_ops.identity_n([1, 1])
+
+  def testMixedTypeListInputFastPathDifferentArity(self):
+    # This tests that the FunctionDef cache key contains the number of args.
+    array_ops.identity_n([self._m_2_by_2, self._m_2_by_2])
+    array_ops.identity_n([self._m_2_by_2, self._m_2_by_2, self._m_2_by_2])
+
+  def testMixedTypeListInputEagerFallbackDifferentArity(self):
+    array_ops.identity_n([1, 1])
+    array_ops.identity_n([1, 1, 1])
+
+  def testSingleTypeListFastPath(self):
+    array_ops.concat([self._m_2_by_2, self._m_2_by_2], axis=-1)
+
+  def testSingleTypeListEagerFallback(self):
+    array_ops.concat([[1], [2]], axis=-1)
+
+  def testSingleTypeListFastPathDifferentArity(self):
+    array_ops.concat([self._m_2_by_2, self._m_2_by_2], axis=-1)
+    array_ops.concat([self._m_2_by_2, self._m_2_by_2, self._m_2_by_2], axis=-1)
+
+  def testSingleTypeListEagerFallbackDifferentArity(self):
+    array_ops.concat([[1], [2]], axis=-1)
+    array_ops.concat([[1], [2], [3]], axis=-1)
+
+  def testCreateCriticalSection(self):
+    cs = critical_section_ops.CriticalSection(shared_name="cs")
+    cs.execute(lambda: 1.0)
+
+
+class RunEagerOpAsFunctionInternalsTest(test.TestCase):
+
+  @test_util.enable_eager_op_as_function
+  def testSimpleGraphExecutesSynchronously(self):
+    if context.num_gpus():
+      self.skipTest("CPU-only test (requires unpartitioned graph).")
+
+    default_executor = test_util.TestDelta("flr_executor", "default")
+    single_threaded = test_util.TestDelta("flr_executor", "single_threaded")
+    run_async = test_util.TestDelta("pflr_runsync", "async")
+    run_sync = test_util.TestDelta("pflr_runsync", "sync")
+    safe = test_util.TestDelta("subgraph_async_summary", "safe_for_sync")
+
+    array_ops.fill([2], constant_op.constant(7, dtype=dtypes.int64))
+
+    assert default_executor.Get() == 0
+    assert single_threaded.Get() > 0
+    assert run_async.Get() == 0
+    assert run_sync.Get() > 0
+    assert safe.Get() > 0
+
+  @test_util.enable_eager_op_as_function
+  def testSendRecvPartitionedGraphExecutesSynchronously(self):
+    if not context.num_gpus():
+      self.skipTest("GPU-only test (requires partitioned graph).")
+
+    default_executor = test_util.TestDelta("flr_executor", "default")
+    single_threaded = test_util.TestDelta("flr_executor", "single_threaded")
+    run_async = test_util.TestDelta("pflr_runsync", "async")
+    run_sync = test_util.TestDelta("pflr_runsync", "sync")
+    send_only = test_util.TestDelta("subgraph_async_summary", "send_only")
+    recv_only = test_util.TestDelta("subgraph_async_summary", "recv_only")
+
+    array_ops.fill([2], constant_op.constant(7, dtype=dtypes.int64))
+
+    assert default_executor.Get() == 0
+    assert single_threaded.Get() > 0
+    assert run_async.Get() == 0
+    assert run_sync.Get() > 0
+    assert send_only.Get() > 0
+    assert recv_only.Get() > 0
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/eager/run_eager_op_as_function_xla_test.py

@@ -0,0 +1,46 @@
+# Copyright 2021 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 wrapping an eager op in a call op at runtime."""
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+
+
+@test_util.with_eager_op_as_function(only_as_function=True)
+class RunEagerOpAsFunctionXlaTest(xla_test.XLATestCase):
+
+  def testVarCreateReadDestroy(self):
+    with self.test_scope():
+      var = variables.Variable(1.0)
+      value = var.read_value()
+      self.assertAllEqual(value, 1.0)
+
+  def testReadVariableInFunction(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      @def_function.function
+      def f():
+        return v.read_value()
+
+      var = f()
+      self.assertEqual(1.0, var.numpy())
+
+
+if __name__ == "__main__":
+  test.main()