chore: import upstream snapshot with attribution

tensorflow/core/function/README.md

@@ -0,0 +1,5 @@
+TensorFlow functions are user-defined callable abstractions built from a graph
+of operations, which represents the function body.
+
+This directory contains function-specific code, spanning from the high-level,
+tf.function API, down to (and excluding) the graph execution framework.

tensorflow/core/function/capture/BUILD

@@ -0,0 +1,104 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+)
+
+py_library(
+    name = "free_vars_detect",
+    srcs = [
+        "free_vars_detect.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/pyct:naming",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:transformer",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+    ],
+)
+
+py_test(
+    name = "free_vars_detect_test",
+    srcs = ["free_vars_detect_test.py"],
+    strict_deps = True,
+    tags = [
+        "no_oss",  # TODO(b/247102978)
+    ],
+    deps = [
+        ":free_vars_detect",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow/python/util:tf_decorator_py",
+    ],
+)
+
+py_test(
+    name = "by_ref_capture_test",
+    srcs = ["by_ref_capture_test.py"],
+    strict_deps = True,
+    deps = [
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+pytype_strict_library(
+    name = "capture_container",
+    srcs = [
+        "capture_container.py",
+    ],
+    deps = [
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python:pywrap_tfe",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:object_identity",
+    ],
+)
+
+py_test(
+    name = "capture_container_test",
+    srcs = ["capture_container_test.py"],
+    strict_deps = True,
+    deps = [
+        ":capture_container",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "restore_captures",
+    srcs = ["restore_captures.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:handle_data_util",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/trackable:asset",
+        "//tensorflow/python/trackable:resource",
+    ],
+)

tensorflow/core/function/capture/by_ref_capture_test.py

@@ -0,0 +1,109 @@
+# 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 detecting free vars in tf.function."""
+
+import unittest
+
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+class ByRefCaptureTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(val_type=[int, constant_op.constant]))
+  def test_direct_capture_mutation(self, val_type):
+    x = val_type(1)
+
+    @def_function.function
+    def f():
+      graph = ops.get_default_graph()
+      cap_x = graph._experimental_capture_side_input_by_ref("x", lambda: x)
+      return cap_x + 1
+
+    self.assertEqual(f(), 2)
+    x = val_type(2)
+    self.assertEqual(f(), 3)
+
+  @unittest.skip("By ref capture API does not work for nested tf.function.")
+  def test_capture_in_nested_function(self):
+    x = constant_op.constant(1)
+
+    @def_function.function
+    def f():
+      graph = ops.get_default_graph()
+      # Capture the same x for the outer tf.function
+      graph._experimental_capture_side_input_by_ref("x", lambda: x)
+
+      @def_function.function
+      def g():
+        graph = ops.get_default_graph()
+        cap_x = graph._experimental_capture_side_input_by_ref("xx", lambda: x)
+        return cap_x + 100
+
+      return g()
+
+    self.assertEqual(f(), 2)
+    x = constant_op.constant(2)
+    self.assertEqual(f(), 102)
+
+  def test_capture_in_outer_function(self):
+    x = 1
+
+    def g():
+      graph = ops.get_default_graph()
+      cap_x = graph._experimental_capture_side_input_by_ref("x", lambda: x)
+      return cap_x + 1
+
+    @def_function.function
+    def f():
+      return g()
+
+    self.assertEqual(f(), 2)
+    x = 2
+    self.assertEqual(f(), 3)
+
+  @unittest.skip("By ref capture API does not work for nested tf.function.")
+  def test_capture_in_outer_tf_function(self):
+    x = 1
+
+    @def_function.function
+    def g():
+      graph = ops.get_default_graph()
+      cap_x = graph._experimental_capture_side_input_by_ref("x", lambda: x)
+      return cap_x + 1
+
+    @def_function.function
+    def f():
+      # Call `_experimental_capture_side_input_by_ref` so that the outer
+      # tf.function will retrace when needed.
+      graph = ops.get_default_graph()
+      graph._experimental_capture_side_input_by_ref("x", lambda: x)
+      return g()
+
+    self.assertEqual(f(), 2)
+    x = 2
+    self.assertEqual(f(), 3)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/core/function/capture/capture_container.py

@@ -0,0 +1,344 @@
+# 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.
+# ==============================================================================
+"""FuncGraph and related functionality."""
+
+import collections as py_collections
+import functools
+from typing import Any, Callable, Hashable, Mapping, Optional
+
+from tensorflow.core.function import trace_type
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.framework import dtypes
+from tensorflow.python.types import core
+from tensorflow.python.util import object_identity
+
+
+_EAGER_CONST_THRESHOLD = 128
+
+
+class MutationAwareDict(py_collections.OrderedDict):
+  """A dict with a mutation flag."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self._mutated = True
+
+  def pop(self, key, default=None):
+    self._mutated = True
+    return super().pop(key, default)
+
+  def __setitem__(self, key, value):
+    self._mutated = True
+    return super().__setitem__(key, value)
+
+  def __delitem__(self, key):
+    self._mutated = True
+    return super().__delitem__(key)
+
+  def clear(self):
+    self._mutated = True
+    return super().clear()
+
+  @property
+  def mutated(self):
+    return self._mutated
+
+  @mutated.setter
+  def mutated(self, value):
+    self._mutated = value
+
+
+class FunctionCaptures(object):
+  """A container for all capture usages within FuncGraph."""
+
+  def __init__(self):
+    self._by_ref_internal = py_collections.OrderedDict()
+    self._by_ref_external = py_collections.OrderedDict()
+    self._by_ref_tracetype = py_collections.OrderedDict()
+    self._by_val_internal = MutationAwareDict()
+    self._by_val_external = MutationAwareDict()
+    self._by_val_tracetype = py_collections.OrderedDict()
+
+    # Set of external ops on which the graph has a control dependency
+    self.control = object_identity.ObjectIdentitySet()
+
+    # Cached properties derived from the above.
+    self._cached_by_val_capture_tuples = []
+    self._cached_capture_types = py_collections.OrderedDict()
+
+  def clear(self):
+    self._by_ref_internal.clear()
+    self._by_ref_external.clear()
+    self._by_ref_tracetype.clear()
+    self._by_val_internal.clear()
+    self._by_val_external.clear()
+
+  def capture_by_value(
+      self, graph: Any, tensor: core.Tensor, name: Optional[str] = None
+  ) -> core.Tensor:
+    """Captures `tensor` if it's external to this graph.
+
+    If `tensor` is from a different graph, returns a placeholder for it.
+    `tensor` and the placeholder will appear in self.captures, and the
+    placeholder will appear in self.inputs.  Multiple calls to this method with
+    the same `tensor` argument will return the same placeholder. If `tensor` is
+    from this graph, returns `tensor`.
+
+    Args:
+      graph: The FuncGraph that captures this tensor.
+      tensor: Tensor. May be from this FuncGraph or a different graph.
+      name: Optional name if a placeholder is created.
+
+    Returns:
+      Tensor from this FuncGraph.
+
+    Raises:
+      InaccessibleTensorError: if any tensors are accessed in a manner that
+      bypasses the mechanisms required for the data dependencies to be correctly
+      wired.
+    """
+    if isinstance(tensor, core.Value):
+      if name is None:
+        # A unique (within the program execution) integer.
+        name = str(pywrap_tfe.TFE_Py_UID())
+
+      # Small EagerTensors are captured with Const ops
+      if (
+          tensor.dtype in dtypes.TF_VALUE_DTYPES
+          and functools.reduce(lambda a, b: a * b, tensor.shape, 1)
+          <= _EAGER_CONST_THRESHOLD
+      ):
+        graph_const = self.by_val_internal.get(id(tensor))
+        if graph_const is None:
+          graph_const = tensor._capture_as_const(name)  # pylint: disable=protected-access
+          if graph_const is None:
+            # Some eager tensors, e.g. parallel tensors, are not convertible to
+            # a single constant. We'll use a placeholder for this case.
+            graph_const = self._create_placeholder_helper(graph, tensor, name)
+          self.add_or_replace(
+              key=id(tensor),
+              external=tensor,
+              internal=graph_const,
+              is_by_ref=False,
+          )
+          graph.inputs.append(graph_const)
+        graph_const._record_tape(tensor)  # pylint: disable=protected-access
+        return graph_const
+
+      # Large EagerTensors and resources are captured with Placeholder ops
+      return self._create_placeholder_helper(graph, tensor, name)
+
+    if tensor.graph is not graph:
+      graph._validate_in_scope(tensor)  # pylint: disable=protected-access
+      if name is None:
+        assert tensor.op is not None, (
+            tensor.__class__,
+            dir(tensor),
+            tensor.__class__.__name__,
+        )
+        name = tensor.op.name
+      # cond/while graphs override _capture_helper() so cannot call
+      # self.create_placeholder_helper() here directly.
+      return graph._capture_helper(tensor, name)  # pylint: disable=protected-access
+    return tensor
+
+  def add_or_replace(
+      self,
+      key: Hashable,
+      external: Any,
+      internal: core.Tensor,
+      tracetype: Any = None,
+      is_by_ref: bool = False,
+  ) -> None:
+    """Replace a already exsiting capture, otherwise add it."""
+    if is_by_ref:
+      self._by_ref_external[key] = external
+      self._by_ref_internal[key] = internal
+      self._by_ref_tracetype[key] = tracetype
+    else:
+      self._by_val_internal[key] = internal
+      self._by_val_external[key] = external
+      if tracetype is not None:
+        self._by_val_tracetype[key] = tracetype
+      else:
+        self._by_val_tracetype[key] = trace_type.from_value(external)
+
+  def pop(self, key: Hashable, is_by_ref: bool = False) -> Any:
+    if is_by_ref:
+      return (
+          self._by_ref_external.pop(key, None),
+          self._by_ref_internal.pop(key, None),
+          self._by_ref_tracetype.pop(key, None),
+      )
+    else:
+      return (
+          self._by_val_external.pop(key, None),
+          self._by_val_internal.pop(key, None),
+          self._by_val_tracetype.pop(key, None),
+      )
+
+  def reset_captures(self, tensors, placeholders):
+    """Set the captures with the provided list of captures & placeholder."""
+    self._by_val_external = MutationAwareDict()
+    self._by_val_internal = MutationAwareDict()
+    self._by_val_tracetype = MutationAwareDict()
+    for external, internal in zip(tensors, placeholders):
+      key = id(external)
+      self._by_val_external[key] = external
+      self._by_val_internal[key] = internal
+      self._by_val_tracetype[key] = trace_type.from_value(external)
+
+  # TODO(panzf): make the method public after supporting lam() returns
+  # non-tensor values. Currently, this method is only used by
+  # FuncGraph._experimental_capture_side_input_by_ref(), which contains the
+  # logics for converting non-tensor values to tensor.
+  def _capture_by_ref(
+      self, graph: Any, lam: Callable[[], Any], key: Hashable = None
+  ) -> Any:
+    """Used during tracing process to create/retrive by-ref captures.
+
+    Args:
+      graph: The FuncGraph that captures this tensor.
+      lam: A callable that takes no arguments and returns tensor captures.
+      key: A hashable identifier.
+
+    Returns:
+      Tensor from this FuncGraph.
+    """
+    # Check if the capture exists in self._by_ref
+    if key is not None and key in self._by_ref_internal:
+      return self._by_ref_internal[key]
+    if key is None:
+      key = len(self._by_ref_internal)
+      while key in self._by_ref_internal:
+        key += 1
+
+    value_nested = lam()
+    capture_trace_type = trace_type.from_value(value_nested)
+    ctx = trace_type.InternalPlaceholderContext(graph)
+    internal = capture_trace_type.placeholder_value(ctx)
+
+    def lam_fn():
+      # pytype: disable=attribute-error
+      value = lam()
+      return capture_trace_type.to_tensors(value)
+      # pytype: enable=attribute-error
+
+    self._by_ref_external[key] = lam_fn
+    self._by_ref_internal[key] = internal
+    self._by_ref_tracetype[key] = capture_trace_type
+    return self._by_ref_internal[key]
+
+  def merge_by_ref_with(self, other: "FunctionCaptures") -> None:
+    """Add by-ref captures from `other` to `self` if not exist."""
+    assert isinstance(other, FunctionCaptures)
+    for key in other.by_ref_external:
+      if key not in self._by_ref_external:
+        self._by_ref_external[key] = other.by_ref_external[key]
+        self._by_ref_tracetype[key] = other.by_ref_tracetype[key]
+
+  # TODO(panzf): Return structured values instead of flat tensors.
+  def get_by_ref_snapshot(self) -> Mapping[Hashable, Any]:
+    """Get a snapshot of current values of by-ref captures."""
+    snapshot = {}
+    for key in self._by_ref_external:
+      func = self._by_ref_external[key]
+      try:
+        value = func()
+      except (AttributeError, RuntimeError):
+        # b/269680071 In case of by-ref captures are unavailable at dispatch
+        # time, use the predefined trace_type instead.
+        value = self._by_ref_tracetype[key]
+      snapshot[key] = value
+    return snapshot
+
+  def _create_placeholder_helper(
+      self, graph: Any, tensor: core.Tensor, name: str
+  ):
+    """A helper function to create capture placeholder."""
+    placeholder = self._by_val_internal.get(id(tensor))
+    if placeholder is None:
+      tracing_ctx = trace_type.InternalTracingContext()
+      spec = trace_type.from_value(tensor, tracing_ctx)
+      spec._name = name  # pylint: disable=protected-access
+      if isinstance(tensor, core.Value) and tensor.is_packed:
+        composite_device_name = tensor.device
+      else:
+        composite_device_name = None
+      placeholder_ctx = trace_type.InternalPlaceholderContext(
+          graph,
+          with_none_control_dependencies=True,
+          composite_device_name=composite_device_name,
+      )
+      placeholder = spec.placeholder_value(placeholder_ctx)
+      self.add_or_replace(
+          key=id(tensor), external=tensor, internal=placeholder, is_by_ref=False
+      )
+      graph.inputs.append(placeholder)
+    placeholder._record_tape(tensor)  # pylint: disable=protected-access
+    return placeholder
+
+  def _recompute_cached_properties(self):
+    """Regenerates cached properties if there have been mutations."""
+    self._by_val_internal.mutated = False
+    self._by_val_external.mutated = False
+    assert len(self._by_val_internal) == len(self._by_val_external)
+    self._cached_by_val_capture_tuples = []
+    for key in self._by_val_internal:
+      assert key in self._by_val_external
+      internal = self._by_val_internal[key]
+      external = self._by_val_external[key]
+      self._cached_by_val_capture_tuples.append((external, internal))
+
+    self._cached_capture_types = py_collections.OrderedDict(
+        list(self._by_val_tracetype.items())
+        + list(self._by_ref_tracetype.items())
+    )
+
+  @property
+  def capture_types(self):
+    if self._by_val_internal.mutated or self._by_val_external.mutated:
+      self._recompute_cached_properties()
+    return self._cached_capture_types
+
+  @property
+  def by_val_capture_tuples(self):
+    if self._by_val_internal.mutated or self._by_val_external.mutated:
+      self._recompute_cached_properties()
+    return self._cached_by_val_capture_tuples
+
+  @property
+  def by_ref_internal(self):
+    return self._by_ref_internal
+
+  @property
+  def by_ref_external(self):
+    return self._by_ref_external
+
+  @property
+  def by_ref_tracetype(self):
+    return self._by_ref_tracetype
+
+  @property
+  def by_val_internal(self):
+    return self._by_val_internal
+
+  @property
+  def by_val_external(self):
+    return self._by_val_external
+
+  @property
+  def by_val_tracetype(self):
+    return self._by_val_tracetype

tensorflow/core/function/capture/capture_container_test.py

@@ -0,0 +1,169 @@
+# 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 tf.function capture containers."""
+import copy
+
+from absl.testing import parameterized
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.capture import capture_container
+from tensorflow.python.platform import test
+
+
+class MutationAwareDictTest(test.TestCase, parameterized.TestCase):
+
+  def _prepare_dict(self):
+    d = {"a": 1, "b": 2, "c": 3}
+    mutation_d = capture_container.MutationAwareDict(copy.copy(d))
+    return d, mutation_d
+
+  @parameterized.parameters(
+      ("__contains__", "a"),
+      ("__contains__", "not_exist"),
+      ("__len__", None),
+      ("__getitem__", "a"))
+  def test_same_behavior_with_normal_dict(self, method, arg):
+    d, mutation_d = self._prepare_dict()
+    d_method = getattr(d, method)
+    mutation_d_method = getattr(mutation_d, method)
+    if arg is None:
+      self.assertEqual(d_method(), mutation_d_method())
+    else:
+      self.assertEqual(d_method(arg), mutation_d_method(arg))
+
+  @parameterized.parameters(
+      ("pop",),
+      ("__delitem__",))
+  def test_pop_and_del(self, method):
+    d, mutation_d = self._prepare_dict()
+    d_method = getattr(d, method)
+    mutation_d_method = getattr(mutation_d, method)
+    d_method("b")
+    mutation_d_method("b")
+    self.assertListEqual(list(d.keys()), list(mutation_d.keys()))
+    self.assertListEqual(list(d.values()), list(mutation_d.values()))
+
+  def test_mutatation_ops(self):
+    _, d = self._prepare_dict()
+    with self.subTest("set"):
+      d["d"] = 4
+      self.assertTrue(d.mutated)
+    with self.subTest("pop"):
+      d.pop("d")
+      self.assertTrue(d.mutated)
+    with self.subTest("del"):
+      del d["c"]
+      self.assertTrue(d.mutated)
+    with self.subTest("clear"):
+      d.clear()
+      self.assertTrue(d.mutated)
+
+  def test_mutated_property(self):
+    _, d = self._prepare_dict()
+    with self.subTest("initial_state"):
+      self.assertTrue(d.mutated)
+    with self.subTest("setter"):
+      d.mutated = False
+      self.assertFalse(d.mutated)
+
+
+class FunctionCapturesTest(test.TestCase, parameterized.TestCase):
+
+  def test_add_or_replace(self):
+    fn_captures = capture_container.FunctionCaptures()
+    fn_captures.add_or_replace("a", 1, -1, is_by_ref=False)
+    fn_captures.add_or_replace("aa", 1, -1, 0, is_by_ref=True)
+
+    with self.subTest("add_by_val"):
+      self.assertLen(fn_captures.by_val_internal, 1)
+      self.assertLen(fn_captures.by_val_external, 1)
+
+    with self.subTest("add_by_ref"):
+      self.assertLen(fn_captures.by_ref_internal, 1)
+      self.assertLen(fn_captures.by_ref_external, 1)
+      self.assertLen(fn_captures.by_ref_tracetype, 1)
+
+    fn_captures.add_or_replace("a", 2, -2, is_by_ref=False)
+    with self.subTest("replace_by_val"):
+      self.assertLen(fn_captures.by_val_internal, 1)
+      self.assertLen(fn_captures.by_val_external, 1)
+      self.assertEqual(fn_captures.by_val_external["a"], 2)
+      self.assertEqual(fn_captures.by_val_internal["a"], -2)
+
+  def test_by_val_capture_tuples(self):
+    fn_captures = capture_container.FunctionCaptures()
+
+    with self.subTest("initial_state"):
+      self.assertEmpty(fn_captures.by_val_capture_tuples)
+
+    with self.subTest("add"):
+      fn_captures.add_or_replace("a", 1, -1, is_by_ref=False)
+      self.assertLen(fn_captures.by_val_capture_tuples, 1)
+      self.assertSequenceEqual(
+          fn_captures.by_val_capture_tuples,
+          ((1, -1),))
+
+      fn_captures.add_or_replace("b", 2, -2, is_by_ref=False)
+      self.assertLen(fn_captures.by_val_capture_tuples, 2)
+      self.assertSequenceEqual(
+          fn_captures.by_val_capture_tuples,
+          ((1, -1), (2, -2)))
+
+    with self.subTest("replace"):
+      fn_captures.add_or_replace("a", 1, -3, is_by_ref=False)
+      self.assertLen(fn_captures.by_val_capture_tuples, 2)
+      self.assertSequenceEqual(
+          fn_captures.by_val_capture_tuples,
+          ((1, -3), (2, -2)))
+
+    with self.subTest("pop"):
+      fn_captures.pop("b", is_by_ref=False)
+      self.assertSequenceEqual(
+          fn_captures.by_val_capture_tuples,
+          ((1, -3),))
+
+    with self.subTest("reset"):
+      fn_captures.reset_captures([10, 20], [-10, -20])
+      self.assertSequenceEqual(
+          fn_captures.by_val_capture_tuples,
+          ((10, -10), (20, -20)))
+
+    with self.subTest("clear"):
+      fn_captures.clear()
+      self.assertEmpty(fn_captures.by_val_capture_tuples)
+
+  def test_capture_types(self):
+    class FakePlaceholder():
+      pass
+
+    fn_captures = capture_container.FunctionCaptures()
+    fn_captures.add_or_replace("v1", 1, FakePlaceholder(), is_by_ref=False)
+    fn_captures.add_or_replace("v2", 2, FakePlaceholder(), is_by_ref=False)
+    fn_captures.add_or_replace("v3", 3, FakePlaceholder(), is_by_ref=False)
+    fn_captures.add_or_replace(
+        "r1", 1, FakePlaceholder(), trace_type.from_value(4), is_by_ref=True)
+    fn_captures.add_or_replace(
+        "r2", 2, FakePlaceholder(), trace_type.from_value(5), is_by_ref=True)
+
+    self.assertLen(fn_captures.capture_types, 5)
+    self.assertEqual(fn_captures.capture_types["v1"], trace_type.from_value(1))
+    self.assertEqual(fn_captures.capture_types["v2"], trace_type.from_value(2))
+    self.assertEqual(fn_captures.capture_types["v3"], trace_type.from_value(3))
+    self.assertEqual(fn_captures.capture_types["r1"], trace_type.from_value(4))
+    self.assertEqual(fn_captures.capture_types["r2"], trace_type.from_value(5))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/core/function/capture/free_vars_detect.py

@@ -0,0 +1,329 @@
+# 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.
+# ==============================================================================
+"""An independent module to detect free vars inside a function."""
+
+import builtins
+import collections
+import functools
+import inspect
+import types
+
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import inspect_utils
+from tensorflow.python.autograph.pyct import naming
+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
+
+FreeVar = collections.namedtuple("FreeVar", ["name", "is_function", "obj"])
+
+_fn_log_cache = dict()
+
+
+def _parse_and_analyze(func):
+  """Parse and analyze Python Function code."""
+  node, source = parser.parse_entity(func, future_features=())
+  node = qual_names.resolve(node)
+  entity_info = transformer.EntityInfo(
+      name=func.__name__,
+      source_code=source,
+      source_file=None,
+      future_features=(),
+      namespace={})
+  namer = naming.Namer({})
+  ctx = transformer.Context(entity_info, namer, None)
+  node = activity.resolve(node, ctx)
+  return node
+
+
+def _handle_wrap_partial_func(obj):
+  """Processes wrapped function and partial functions recursively."""
+  modified = True
+  while modified:
+    modified = False
+    while hasattr(obj, "__wrapped__"):
+      obj = obj.__wrapped__
+      modified = True
+    if isinstance(obj, functools.partial) or isinstance(
+        obj, functools.partialmethod):
+      obj = obj.func
+      modified = True
+  return obj
+
+
+def _get_self_obj_from_closure(fn):
+  """Get the object that `self` keyword refers to within a function.
+
+  Args:
+    fn: A python function object
+
+  Returns:
+    A class object that `self` refers to. Return None if not found.
+
+  Here is an example demonstrating how this helper function works.
+
+  ```
+  class Foo():
+
+    def __init__(self):
+      self.val = 1
+
+    def bar(self):
+      x = 2
+
+      def fn():
+        return self.val + x
+
+      return fn
+
+  foo = Foo()
+  fn = foo.bar()
+  self_obj = _get_self_obj_from_closure(fn)
+  assert self_obj is foo
+  ```
+
+  The goal is to get the `self_obj` (foo) from `fn`, so that it's feasible to
+  access attributes of `foo`, like self.val in this case.
+
+  This function first parses fn.qual_name, "Foo.bar.<locals>.fn", and finds the
+  closure whose class name appear in fn.qual_name first.
+  """
+  assert hasattr(fn, "__closure__")
+  qual_name = fn.__qualname__.split(".")
+  # Search from the right to left
+  qual_name = qual_name[::-1]
+
+  if fn.__closure__:
+    for cls_name in qual_name:
+      for cell in fn.__closure__:
+        try:
+          closure = cell.cell_contents
+        except ValueError:
+          # Continue when cell is empty and its content is unavailable
+          continue
+        if inspect.isclass(type(closure)):
+          if type(closure).__name__ == cls_name:
+            obj = closure
+            return obj
+
+  return None
+
+
+def _search_callable_free_vars(fn):
+  """Search free vars from a callable object."""
+  fn = _handle_wrap_partial_func(fn)
+
+  try:
+    node = _parse_and_analyze(fn)
+  except ValueError:
+    # When source code unavailable, return empty result
+    return []
+  except NotImplementedError:
+    # Autograph cannot handle multiple lambda functions with same line number
+    # and args name.
+    return []
+
+  scope = anno.getanno(node, anno.Static.SCOPE)
+  free_vars_all = list(scope.free_vars)
+  namespace = inspect_utils.getnamespace(fn)
+  filtered = []
+
+  for var in free_vars_all:
+    base = str(var.qn[0])
+
+    if var.is_simple():
+      if base in builtins.__dict__.keys():
+        continue
+      obj = namespace.get(base, None)
+    else:
+      assert var.is_composite()
+      # A compositve qualified name `QN` can be either an attr or a subscript
+      if var.has_subscript():
+        # For free var with subscripts, both the base and full formats are
+        # generated.
+        # For example, if the code have `glob[idx]`, `free_vars_all` would
+        # contain `glob` as well as `glob[idx]`.
+        # The method only keeps the base format for simplicity.
+        continue
+      else:
+        assert var.has_attr()
+        # For free vars with multiple attributes like `f.g.h`,
+        # just as the subscripts, multiple free vars (QN) are generated:
+        # ['f', 'f.g', 'f.g.h']
+        # If `f` is `self`, only process the first attribute `f.g`.
+        # Otherwise, only process `f`.
+        if not var.qn[0].is_composite() and base == "self":
+          attr = str(var.qn[1])
+          # For method, access the object that `self` refers to via __self__
+          if hasattr(fn, "__self__"):
+            obj = getattr(fn.__self__, attr, None)
+          # For function (not method) `self` usage under enclosing class scope
+          elif hasattr(fn, "__closure__"):
+            self_obj = _get_self_obj_from_closure(fn)
+            if self_obj:
+              obj = getattr(self_obj, attr, None)
+            else:
+              continue
+          else:
+            continue
+        else:
+          continue
+
+    if (inspect.ismodule(obj) or inspect.isclass(obj)):
+      continue
+    elif inspect.isfunction(obj) or inspect.ismethod(obj):
+      obj = _handle_wrap_partial_func(obj)
+      if obj.__module__ != fn.__module__:
+        continue
+      filtered.append(FreeVar(str(var), True, obj))
+    else:
+      filtered.append(FreeVar(str(var), False, None))
+
+  filtered = sorted(filtered, key=lambda x: x.name)
+  return filtered
+
+
+def _make_lambda_name(obj):
+  source = inspect.getsource(obj)
+  name = source.split("=")[0].strip()
+  return name
+
+
+def _make_callable_signature(obj):
+  """Generate signature for function/method."""
+  if inspect.isclass(obj) or inspect.isfunction(obj):
+    if obj.__name__ == "<lambda>":
+      return _make_lambda_name(obj)
+    return obj.__name__
+  elif inspect.ismethod(obj):
+    obj_self = obj.__self__
+    if isinstance(obj_self, type):
+      cls_name = obj_self.__name__
+    else:
+      cls_name = obj_self.__class__.__name__
+    return f"{cls_name}.{obj.__name__}"
+  else:
+    raise TypeError(
+        f"Only class/function/methods are valid inputs, got {type(obj)}")
+
+
+def _detect_function_free_vars(fn):
+  """Detect free vars in any Python function."""
+  assert isinstance(fn, types.FunctionType) or isinstance(
+      fn, types.MethodType
+  ), f"The input should be of Python function type. Got type: {type(fn)}."
+
+  queue = collections.deque([fn])
+  fn_map = dict()
+
+  # Perform BFS over functions to get free vars
+  while queue:
+    obj = queue.popleft()
+    signature = _make_callable_signature(obj)
+    if signature not in fn_map:
+      free_vars = _search_callable_free_vars(obj)
+      if not free_vars:
+        continue
+      fn_map[signature] = free_vars
+      for var in free_vars:
+        # Only search callable free vars
+        if var.is_function:
+          obj = var.obj
+          if _make_callable_signature(obj) not in fn_map:
+            queue.append(obj)
+
+  # func_name -> namedtupe FreeVar
+  return fn_map
+
+
+def generate_free_var_logging(fn, fn_threshold=5, var_threshold=10):
+  """Generate loggings of free vars from fn."""
+  # Now only detect free vars for function/method
+  if not (isinstance(fn, types.FunctionType) or isinstance(
+      fn, types.MethodType) or isinstance(fn, functools.partial) or
+          isinstance(fn, functools.partialmethod)):
+    return None
+  fn = _handle_wrap_partial_func(fn)
+
+  if not (hasattr(fn, "__module__") and hasattr(fn, "__qualname__")):
+    return None
+  fn_key = (fn.__module__, fn.__qualname__)
+  # To prevent log spam, only generate logging once for each tf.function
+  if fn_key in _fn_log_cache:
+    return None
+
+  try:
+    fn_vars_map = _detect_function_free_vars(fn)
+  except Exception:  # pylint: disable=broad-except
+    # Only for logging purpose, do not raise errors to users
+    return None
+
+  # If not free vars detected, return None
+  if not fn_vars_map:
+    _fn_log_cache[fn_key] = None
+    return _fn_log_cache[fn_key]
+
+  logging_txt = []
+  tf_fn_name = _make_callable_signature(fn)
+  tf_fn_module = fn.__module__
+
+  def one_line_logging(fn_name, free_vars, threshold=10):
+    if not free_vars:
+      return ""
+    log = f"Inside function {fn_name}(): "
+    log += ", ".join([var.name for var in free_vars[:threshold]])
+    if len(free_vars) > threshold:
+      log += "..."
+    return log
+
+  # Show the free vars info of the tf.function at the top
+  fn_threshold -= 1
+  try:
+    tf_fn_line = one_line_logging(tf_fn_name, fn_vars_map[tf_fn_name],
+                                  var_threshold)
+  except Exception:  # pylint: disable=broad-except
+    # Only for logging purpose, do not raise error to users
+    return ""
+
+  # Functions that are defined outside of tf.function
+  outside_fn_lines = []
+  outside_fn_names = [name for name in fn_vars_map.keys() if name != tf_fn_name]
+  outside_fn_names = sorted(outside_fn_names)
+  for fn_name in outside_fn_names[:fn_threshold]:
+    outside_fn_lines.append(
+        one_line_logging(fn_name, fn_vars_map[fn_name], var_threshold))
+
+  if len(fn_vars_map) > fn_threshold:
+    ellipsis_line = "..."
+  else:
+    ellipsis_line = None
+
+  # TODO(panzf): direct users to the manual API after it's exposed to public
+  explanation_line = (
+      f"Free variables are detected within tf.function {tf_fn_name}() in"
+      f"{tf_fn_module}. Free variable usage may cause inconsistant behaviors"
+      "between eager mode and tf.function. Please consider refactor the code"
+      "if possible. More details are avaiable in"
+      "https://www.tensorflow.org/guide/function#limitations.\n"
+      "Free variable names inside each function/method are shown below:")
+
+  logging_txt = [explanation_line, tf_fn_line] + outside_fn_lines
+  if ellipsis_line:
+    logging_txt.append(ellipsis_line)
+  logging_txt = "\n".join(logging_txt)
+
+  _fn_log_cache[fn_key] = logging_txt
+  return _fn_log_cache[fn_key]

tensorflow/core/function/capture/free_vars_detect_test.py

@@ -0,0 +1,736 @@
+# 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 detecting free vars in tf.function."""
+
+import functools
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.function.capture import free_vars_detect
+from tensorflow.python.util import tf_decorator
+
+
+def get_var_name(d):
+  return [var.name for var in d]
+
+
+class GetSelfObjFromClosureTest(parameterized.TestCase):
+
+  def test_single_enclosing_class(self):
+
+    class Foo():
+
+      def __init__(self):
+        self.val = 1
+
+      def bar(self):
+        x = 2
+
+        def fn():
+          return self.val + x
+
+        return fn
+
+    foo = Foo()
+    fn = foo.bar()
+    self_obj = free_vars_detect._get_self_obj_from_closure(fn)
+    self.assertIs(self_obj, foo)
+
+
+class FreeVarDetectionTest(parameterized.TestCase):
+
+  def test_func_arg(self):
+    x = 1  # pylint: disable=unused-variable
+
+    def f(x):
+      return x + 1
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertEmpty(func_map)
+
+  def test_func_local_var(self):
+
+    def f():
+      x = 1
+      return x + 1
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertEmpty(func_map)
+
+  def test_global_var_int(self):
+    x = 1
+
+    def f():
+      return x + 1
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_builtin_func(self):
+
+    def f(x):
+      return len(x)
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertEmpty(func_map)
+
+  def test_global_var_dict(self):
+    glob = {"a": 1}
+
+    def f():
+      return glob["a"] + 1
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_global_var_dict_w_var_index(self):
+    glob = {"a": 1}
+    key = "a"
+
+    def f():
+      return glob[key] + 1
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["glob", "key"])
+
+  def test_duplicate_global_var(self):
+    x = 1
+
+    def f():
+      return x + x
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  @parameterized.named_parameters(
+      ("lambda_1", lambda _x: 3,), ("lambda_2", lambda _x: 3,))
+  def test_multiple_lambda_w_same_line_num_and_args(self, fn):
+    func_map = free_vars_detect._detect_function_free_vars(fn)
+    self.assertEmpty(func_map)
+
+  def test_lambda_wo_free_var(self):
+    f = lambda x: x + x
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertEmpty(func_map)
+
+  def test_lambda_w_free_var(self):
+    glob = 1
+    f = lambda x: x + glob
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_multi_lambda_w_free_var(self):
+    glob = 1
+    g = lambda x: x + glob
+    h = lambda: glob + 1
+
+    def f(x):
+      return g(x) + h()
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertLen(func_map, 3)
+    self.assertIn("f", func_map.keys())
+    self.assertIn("g", func_map.keys())
+    self.assertIn("h", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["g", "h"])
+    free_vars = get_var_name(func_map["g"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+    free_vars = get_var_name(func_map["h"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_lambda_inline(self):
+    glob = 1
+
+    def f(x):
+      return lambda: x + glob
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_glob_numpy_var(self):
+    a = 0
+    b = np.asarray(1)
+
+    def f():
+      c = np.asarray(2)
+      res = a + b + c
+      return res
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["a", "b"])
+
+  def test_global_var_in_nested_func(self):
+    x = 1
+
+    def f():
+
+      def g():
+        return x + 1
+
+      return g()
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    self.assertLen(func_map.keys(), 1)
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_global_var_from_outer_func(self):
+    x = 1
+
+    def g():
+      return x + 1
+
+    def f():
+      return g()
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    self.assertIn("g", func_map.keys())
+    self.assertLen(func_map.keys(), 2)
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["g"])
+    free_vars = get_var_name(func_map["g"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_method(self):
+    x = 1
+
+    class Foo():
+
+      def f(self):
+        return x
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.f)
+    self.assertLen(func_map.keys(), 1)
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_method_w_method_call(self):
+    x = 0
+
+    class Foo():
+
+      def f(self):
+        return self.g
+
+      def g(self):
+        return [x]
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.f)
+    self.assertLen(func_map.keys(), 2)
+
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["self.g"])
+
+    self.assertIn("Foo.g", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.g"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_method_w_self_as_arg(self):
+    x = 1
+
+    class Foo():
+
+      def f(self):
+        return self.g(self)
+
+      def g(self, obj):
+        if obj != self:
+          return x
+        else:
+          return -x
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.f)
+    self.assertLen(func_map.keys(), 2)
+
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["self.g"])
+
+    self.assertIn("Foo.g", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.g"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_self_inside_method(self):
+    x = 1
+
+    class Foo():
+
+      def __init__(self):
+        self.val = 2
+
+      def bar(self):
+
+        def tf_func():
+          return self.val + x
+
+        return tf_func
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.bar())
+    self.assertLen(func_map.keys(), 1)
+
+    self.assertIn("tf_func", func_map.keys())
+    free_vars = get_var_name(func_map["tf_func"])
+    self.assertSequenceEqual(free_vars, ["self", "self.val", "x"])
+
+  def test_self_inside_function_w_multiple_closures(self):
+    # Test when a function contins multiple closures
+
+    class Foo():
+
+      def method(self):
+
+        class Baz():
+
+          def baz_str(self):
+            return "Baz"
+
+        baz = Baz()
+        x = "x"
+
+        class Bar():
+
+          def bar_str(self):
+            return x + "Bar"
+
+          def method(self):
+
+            def fn():
+              return self.bar_str() + baz.baz_str()
+
+            return fn
+
+        bar = Bar()
+        return bar.method()
+
+    foo = Foo()
+    fn = foo.method()
+    # cells for `self.bar_str()`, `baz.baz_str()`
+    self.assertLen(fn.__closure__, 2)
+
+    func_map = free_vars_detect._detect_function_free_vars(fn)
+    self.assertLen(func_map.keys(), 2)
+
+    self.assertIn("fn", func_map.keys())
+    free_vars = get_var_name(func_map["fn"])
+    self.assertSequenceEqual(free_vars, ["baz", "self", "self.bar_str"])
+
+    self.assertIn("Bar.bar_str", func_map.keys())
+    free_vars = get_var_name(func_map["Bar.bar_str"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_method_w_self_attribute(self):
+    x = 0
+
+    class Foo():
+
+      def __init__(self):
+        self.x = 1
+        self.y = 2
+
+      def f(self):
+        return self.g + self.x + self.y
+
+      def g(self):
+        return x
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.f)
+    self.assertLen(func_map.keys(), 2)
+
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["self.g", "self.x", "self.y"])
+
+    self.assertIn("Foo.g", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.g"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_method_w_multiple_attributes(self):
+    glob = "dummy_value"
+
+    class Foo():
+
+      def f(self):
+        return self.g.h.x.y.z
+
+      def g(self):
+        return glob
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.f)
+    self.assertLen(func_map.keys(), 2)
+
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["self.g"])
+
+    self.assertIn("Foo.g", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.g"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_classmethod_decorator(self):
+    glob = 1
+
+    class Foo():
+
+      @classmethod
+      def f(cls):
+        return glob
+
+    func_map = free_vars_detect._detect_function_free_vars(Foo.f)
+    self.assertLen(func_map.keys(), 1)
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_method_call_classmethod(self):
+    glob = 1
+
+    class Foo():
+
+      def f(self):
+        return self.g()
+
+      @classmethod
+      def g(cls):
+        return glob
+
+    foo = Foo()
+    func_map = free_vars_detect._detect_function_free_vars(foo.f)
+    self.assertLen(func_map.keys(), 2)
+
+    self.assertIn("Foo.f", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.f"])
+    self.assertSequenceEqual(free_vars, ["self.g"])
+
+    self.assertIn("Foo.g", func_map.keys())
+    free_vars = get_var_name(func_map["Foo.g"])
+    self.assertSequenceEqual(free_vars, ["glob"])
+
+  def test_global_var_from_renamed_outer_func(self):
+    x = 1
+
+    def g():
+      return x + 1
+
+    def f():
+      h = g
+      return h()
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    self.assertIn("g", func_map.keys())
+    self.assertLen(func_map.keys(), 2)
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["g"])
+    free_vars = get_var_name(func_map["g"])
+    self.assertSequenceEqual(free_vars, ["x"])
+
+  def test_decorated_method_w_self_no_exception(self):
+    """Test this pattern does not raise any exceptions."""
+
+    def dummy_tf_function(func):
+
+      func_map = free_vars_detect._detect_function_free_vars(func)
+      self.assertLen(func_map, 1)
+      self.assertIn("foo", func_map.keys())
+      free_vars = get_var_name(func_map["foo"])
+      self.assertSequenceEqual(free_vars, ["dummy_tf_function"])
+
+      def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+      return wrapper
+
+    glob = 1
+
+    # This pattern is not fully supported yet in the sense that `self.bar()` is
+    # not inspected so `glob` cannot be detected.
+    # The reason is the neither `self` nor `self.bar` is accessible from the
+    # perspective of dummy_tf_function decorator.
+    # One possible solution is parsing the source code of the whole module,
+    # instead of single function. And probably get the source of `self.bar`
+    # from the AST of the module where `Foo` is defined. One potentail challenge
+    # of this approach is how to locate the decorated function in the AST.
+    class Foo():
+
+      @dummy_tf_function
+      def foo(self):
+        return self.bar()
+
+      def bar(self):
+        return glob
+
+    _ = Foo()
+
+  # Use `wrapper_first` to control different arguments order
+  @parameterized.parameters(
+      (functools.update_wrapper, True),
+      (tf_decorator.make_decorator, False))
+  def test_func_w_decorator(self, make_decorator, wrapper_first):
+    x = 1
+
+    def decorator_foo(func):
+
+      def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+      if wrapper_first:
+        return make_decorator(wrapper, func)
+      else:
+        return make_decorator(func, wrapper)
+
+    @decorator_foo
+    @decorator_foo
+    def f():
+
+      @decorator_foo
+      @decorator_foo
+      def g():
+        return x + 1
+
+      return g()
+
+    func_map = free_vars_detect._detect_function_free_vars(f)
+    self.assertIn("f", func_map.keys())
+    self.assertLen(func_map.keys(), 2)
+    free_vars = get_var_name(func_map["f"])
+    self.assertSequenceEqual(free_vars, ["decorator_foo", "x"])
+
+  # TODO(panzf): test the pattern when callable function args are supported
+  @unittest.skip("Feature not implemented")
+  def test_global_var_from_arg_func(self):
+    x = 1
+
+    def g():
+      return x + 1
+
+    def f(h):
+      return h()
+
+    _ = f(g)
+
+
+class GenerateLoggingTest(parameterized.TestCase):
+
+  def _remove_explanation(self, logging_txt):
+    free_vars = logging_txt.split("\n")
+    self.assertGreater(len(free_vars), 2)
+    return "\n".join(free_vars[2:])
+
+  def test_none_input(self):
+    txt = free_vars_detect.generate_free_var_logging(None)
+    self.assertIsNone(txt)
+
+  def test_non_function_input(self):
+    x = 1
+
+    class Foo():
+
+      def bar(self):
+        return x
+
+    foo = Foo()
+    txt = free_vars_detect.generate_free_var_logging(foo)
+    self.assertIsNone(txt)
+
+  def test_func_wo_source_code(self):
+    code = "def f_exec():\n  return 1"
+    # Use `exec` to generate a function without source code
+    exec(code, globals())  # pylint: disable=exec-used
+    txt = free_vars_detect.generate_free_var_logging(f_exec)  # pylint: disable=undefined-variable
+    self.assertIsNone(txt)
+
+  def test_no_free_var(self):
+
+    def f(x):
+      return x + 1
+
+    txt = free_vars_detect.generate_free_var_logging(f)
+    self.assertIsNone(txt)
+
+  def test_single_func(self):
+    x = 1
+    y = 2
+
+    def f(a):
+      return a + x + y
+
+    txt = free_vars_detect.generate_free_var_logging(f)
+    txt = self._remove_explanation(txt)
+    self.assertEqual(txt, "Inside function f(): x, y")
+
+  def test_nested_func(self):
+    x = 1
+    y = 2
+
+    def g():
+      return y
+
+    def f():
+      return g() + x
+
+    txt = free_vars_detect.generate_free_var_logging(f)
+    txt = self._remove_explanation(txt)
+    lines = txt.split("\n")
+    self.assertLen(lines, 2)
+    self.assertEqual(lines[0], "Inside function f(): g, x")
+    self.assertEqual(lines[1], "Inside function g(): y")
+
+  def test_method_w_method_call(self):
+    x = 0
+
+    class Foo():
+
+      def f(self):
+        return self.g
+
+      def g(self):
+        return [x]
+
+    foo = Foo()
+    txt = free_vars_detect.generate_free_var_logging(foo.f)
+    txt = self._remove_explanation(txt)
+    lines = txt.split("\n")
+    self.assertLen(lines, 2)
+    self.assertEqual(lines[0], "Inside function Foo.f(): self.g")
+    self.assertEqual(lines[1], "Inside function Foo.g(): x")
+
+  def test_partial_func(self):
+    x = 1
+    y = 2
+
+    def f(a):
+      return a + x + y
+
+    partial_f = functools.partial(f, a=0)
+
+    txt = free_vars_detect.generate_free_var_logging(partial_f)
+    txt = self._remove_explanation(txt)
+    self.assertEqual(txt, "Inside function f(): x, y")
+
+  def test_partial_method(self):
+    x = 0
+
+    class Foo():
+
+      def f(self):
+        return self.g
+
+      def g(self):
+        return [x]
+
+      partial_f = functools.partialmethod(f)
+
+    foo = Foo()
+    txt = free_vars_detect.generate_free_var_logging(foo.partial_f)
+    txt = self._remove_explanation(txt)
+    lines = txt.split("\n")
+    self.assertLen(lines, 2)
+    self.assertEqual(lines[0], "Inside function Foo.f(): self.g")
+    self.assertEqual(lines[1], "Inside function Foo.g(): x")
+
+  def test_partial_wrapped_partial_func(self):
+
+    def decorator_foo(func):
+
+      def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+      return functools.update_wrapper(wrapper, func)
+
+    x = 1
+    y = 2
+
+    def f(a, b):
+      return a + b + x + y
+
+    f = functools.partial(f, a=0)
+    f = decorator_foo(f)
+    f = functools.partial(f, b=0)
+
+    txt = free_vars_detect.generate_free_var_logging(f)
+    txt = self._remove_explanation(txt)
+    self.assertEqual(txt, "Inside function f(): x, y")
+
+  def test_freevar_threshold(self):
+    a = b = c = d = e = 1
+
+    def f():
+      return a + b + c + d + e
+
+    txt = free_vars_detect.generate_free_var_logging(f, var_threshold=3)
+    txt = self._remove_explanation(txt)
+    self.assertEqual(txt, "Inside function f(): a, b, c...")
+
+  def test_func_threshold(self):
+    x = 1
+
+    def g():
+      return x
+
+    def h():
+      return x
+
+    def f():
+      return g() + h()
+
+    txt = free_vars_detect.generate_free_var_logging(f, fn_threshold=2)
+    txt = self._remove_explanation(txt)
+    lines = txt.split("\n")
+    self.assertLen(lines, 3)
+    self.assertEqual(lines[0], "Inside function f(): g, h")
+    self.assertEqual(lines[1], "Inside function g(): x")
+    self.assertEqual(lines[2], "...")
+
+  def test_func_second_call_return_none(self):
+    x = 1
+
+    def f():
+      return x
+
+    logging_txt = free_vars_detect.generate_free_var_logging(f)
+    self.assertIsNotNone(logging_txt)
+    logging_txt = free_vars_detect.generate_free_var_logging(f)
+    self.assertIsNone(logging_txt)
+
+
+if __name__ == "__main__":
+  unittest.main()

tensorflow/core/function/capture/restore_captures.py

@@ -0,0 +1,142 @@
+# 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
+"""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 warnings
+
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import handle_data_util
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.trackable import asset
+from tensorflow.python.trackable import resource
+
+
+def get_tensor_from_node(node):
+  """Resolves a saved model graph node into a tensor to be captured.
+
+  Args:
+    node: a tensor, variable, or resource to be resolved into a capturable
+      tensor
+
+  Returns:
+    A list of tensors.
+  Raises:
+    ValueError: if the node cannot be converted into a tensor.
+  """
+  with ops.init_scope():
+    # TODO(b/210144904): Use __tf_tensor__ instead of `is_[...]` checks
+    if getattr(node, "is_distributed_variable", False):
+      return node
+    elif getattr(node, "is_distributed_table", False):
+      return node
+    elif getattr(node, "is_sharded_variable", False):
+      return node
+    elif resource_variable_ops.is_resource_variable(node):
+      return node.handle
+    elif isinstance(node, asset.Asset):
+      return node.asset_path
+    elif tensor_util.is_tf_type(node):
+      return node
+    elif isinstance(node, resource.CapturableResource):
+      # Note: this executes restored functions in the CapturableResource.
+      return node.resource_handle
+    raise ValueError(f"Cannot convert node {node} to tensor.")
+
+
+def restore_captures(concrete_function, inputs):
+  """Restore captures for the concrete function.
+
+  Used at deserialization time.  For functions that are being deserialized,
+  saved model restores objects that tensors were captured from, but functions
+  only know about their tensors -- object information is destroyed by tracing.
+  This additional logic extracts the tensors which the function originally
+  captured.
+
+  Args:
+    concrete_function: the concrete function for which to restore captures
+    inputs: a list tensors or other Python objects (such as variables) which
+      contain tensors that were originally captured by the function
+  """
+  bound_inputs = [get_tensor_from_node(obj) for obj in inputs]
+  # pylint: disable=g-complex-comprehension
+  bound_variables = [
+      obj
+      for obj in inputs
+      if isinstance(
+          obj,
+          (variables_lib.Variable, resource_variable_ops.BaseResourceVariable),
+      )
+  ]
+  # TODO(b/205010575): This is only injecting the captured inputs into the
+  # concrete function, note that we did not modify the FuncGraph
+  # itself.
+  captured_inputs_list = []
+  concrete_function.set_variables(bound_variables)
+  if bound_inputs:
+    for bound_input, internal_capture in zip(
+        bound_inputs, concrete_function.inputs[-len(bound_inputs) :]
+    ):
+      # Distributed inputs have special logic for capturing, so we call their
+      # custom restoration methods
+      if hasattr(bound_input, "__tf_experimental_restore_capture__"):
+        captured_inputs_list.append(
+            bound_input.__tf_experimental_restore_capture__(
+                concrete_function, internal_capture
+            )
+        )
+      else:
+        captured_inputs_list.append(bound_input)
+        concrete_function.graph.replace_capture(bound_input, internal_capture)
+        if internal_capture.dtype == dtypes.resource:
+          if resource_variable_ops.is_resource_variable(bound_input):
+            try:
+              handle = bound_input.handle
+            except ValueError:
+              # For mirrored variables we'll copy handle data for components
+              # as they get captured.
+              pass
+            else:
+              handle_data_util.copy_handle_data(handle, internal_capture)
+          else:
+            # TODO(b/213451747): Remove need to call copy_handle_data
+            handle_data_util.copy_handle_data(bound_input, internal_capture)
+        # Setting "captures" first means "capture" won't create a new
+        # placeholder for this input.
+        concrete_function.graph.capture(bound_input)
+
+  if any([inp is None for inp in captured_inputs_list]):
+    warnings.warn(
+        "Trying to load ShardedVariables using tf.saved_model.load. "
+        "This won't work if using a tf.distribute.Strategy, and may "
+        "use excess memory if not using a Strategy. Ignore this "
+        "warning if using tf.keras.models.load_model."
+    )
+  concrete_function.set_external_captures(captured_inputs_list)
+
+  # Update FunctionType with new captures.
+  if concrete_function.function_type:
+    concrete_function._function_type = function_type_lib.FunctionType(  # pylint: disable=protected-access
+        concrete_function.function_type.parameters.values(),
+        concrete_function.graph.function_captures.capture_types,
+        return_annotation=concrete_function.function_type.output,
+    )

tensorflow/core/function/integration_test/BUILD

@@ -0,0 +1,30 @@
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/core/function/runtime_client:__subpackages__"],
+)
+
+licenses(["notice"])
+
+py_test(
+    name = "side_inputs_test",
+    srcs = ["side_inputs_test.py"],
+    strict_deps = True,
+    deps = [
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_test(
+    name = "side_inputs_manual_api_test",
+    srcs = ["side_inputs_manual_api_test.py"],
+    strict_deps = True,
+    deps = [
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//tensorflow:tensorflow_py",
+    ],
+)

tensorflow/core/function/integration_test/side_inputs_manual_api_test.py

@@ -0,0 +1,235 @@
+# 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 manual API to support side inputs in tf.function."""
+
+import unittest
+from absl.testing import parameterized
+
+import tensorflow as tf
+
+
+# TODO(panzf): Enable tests after the side inputs manual API is supported
+class SideInputsTest(parameterized.TestCase):
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(
+      (1, tf.constant, 2, tf.constant),
+      (1.0, tf.constant, 2.0, tf.constant),
+      (1, int, 2, int),
+      (1.0, float, 2.0, float),
+      (1, int, 2, tf.constant),
+      (1, tf.constant, 2, int))
+  def test_direct_capture(self, val_before, type_before, val_after, type_after):
+    def f():
+      return tf.func.experimental.capture(lambda: x) + 1
+
+    tf_f = tf.function(f)
+    x = type_before(val_before)
+    self.assertEqual(f(), tf_f())
+    x = type_after(val_after)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  def test_direct_capture_mutation(self):
+    def f():
+      cglob = tf.func.experimental.capture(lambda: glob)
+      return cglob[-1] + tf.constant(0)
+
+    tf_f = tf.function(f)
+    glob = [tf.constant(1), tf.constant(2)]
+    self.assertEqual(f(), tf_f())
+    glob.append(tf.constant(3))
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(
+      tf.constant,
+      int)
+  def test_dict_capture_mutation_with_tensor_and_non_tensor(self, capture_type):
+    def f():
+      cd = tf.func.experimental.capture(lambda: d)
+      return cd["val"]
+
+    tf_f = tf.function(f)
+    d = {"int": 1, "tensor": tf.constant(2), "val": capture_type(3)}
+    self.assertEqual(f(), tf_f())
+    d["val"] = capture_type(4)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(tf.constant, int)
+  def test_capture_with_duplicate_usage(self, capture_type):
+    def f():
+      cx = tf.func.experimental.capture(lambda: x)
+      return cx + cx  # should capture x just once.
+
+    tf_f = tf.function(f)
+
+    x = capture_type(1)  # pylint: disable=unused-variable
+    self.assertEqual(f(), tf_f())
+    self.assertLen(tf_f._variable_creation_config._captures_container, 1)
+
+  @unittest.skip("Feature not implemented")
+  def test_local_capture(self):
+    def f():
+      x = tf.constant(0)
+      def g():
+        return tf.func.experimental.capture(lambda: x)
+      return g()
+
+    tf_f = tf.function(f)
+    x = tf.constant(100)  # pylint: disable=unused-variable
+    # a = f()
+    a = 100
+    b = tf_f()
+    self.assertEqual(a, b)
+    x = tf.constant(200)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(
+      tf.constant,
+      int)
+  def test_capture_by_nested_function(self, capture_type):
+    def f():
+      def g():
+        cx = tf.func.experimental.capture(lambda: x)
+        return cx
+      return g()
+
+    tf_f = tf.function(f)
+    x = capture_type(1)  # pylint: disable=unused-variable
+    self.assertEqual(f(), tf_f())
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(tf.constant, int)
+  def test_outer_capture_with_function_call(self, capture_type):
+    def g():
+      cx = tf.func.experimental.capture(lambda: x)
+      return cx
+
+    def f():
+      return g()
+    tf_f = tf.function(f)
+
+    x = capture_type(1)  # pylint: disable=unused-variable
+    self.assertEqual(f(), tf_f())
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(tf.constant, int)
+  def test_outer_capture_with_nested_function_call(self, capture_type):
+    x = capture_type(1)  # pylint: disable=unused-variable
+    def g_factory():
+      def g():
+        cx = tf.func.experimental.capture(lambda: x)
+        return cx
+      return g()
+
+    def f():
+      h = g_factory
+      return h()
+    tf_f = tf.function(f)
+
+    self.assertEqual(f(), tf_f())
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(tf.constant, int)
+  def test_capture_within_function_argument(self, capture_type):
+    def g():
+      cx = tf.func.experimental.capture(lambda: x)
+      return cx
+
+    def f(h):
+      return h()
+    tf_f = tf.function(f)
+
+    x = capture_type(1)  # pylint: disable=unused-variable
+    self.assertEqual(f(g), tf_f(g))
+    x = capture_type(2)
+    self.assertEqual(f(g), tf_f(g))
+
+  @unittest.skip("Feature not implemented")
+  def test_inner_nested_tf_function_raise_error(self):
+    @tf.function
+    def tf_f():
+
+      @tf.function
+      def tf_g():
+        cx = tf.func.experimental.capture(lambda: x)
+        return cx
+
+      return tf_g()
+
+    x = tf.constant(0)  # pylint: disable=unused-variable
+    with self.assertRaisesRegex(
+        NotImplementedError, "Manual side input usage for inner nested"):
+      tf_f()
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(
+      tf.constant,
+      int)
+  def test_outer_nested_tf_function_with_global_capture(self, capture_type):
+    @tf.function
+    def tf_f():
+
+      @tf.function
+      def tf_g(x):
+        return x
+
+      cx = tf.func.experimental.capture(lambda: x)
+      return tf_g(cx)
+
+    x = capture_type(0)  # pylint: disable=unused-variable
+    self.assertEqual(tf_f(), tf.constant(0))
+    x = capture_type(1)
+    self.assertEqual(tf_f(), tf.constant(1))
+
+  @unittest.skip("Feature not implemented")
+  def test_non_callable_function_raise_error(self):
+    def f():
+      return tf.func.experimental.capture(x) + 1
+
+    tf_f = tf.function(f)
+    x = 1
+    with self.assertRaises(TypeError):
+      _ = tf_f()
+    x = tf.constant(1)
+    with self.assertRaises(TypeError):
+      _ = tf_f()
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(
+      (1, tf.constant, 2, tf.constant),
+      (1, int, 2, int))
+  def test_call_by_value(self, val_before, type_before, val_after, type_after):
+    def f():
+      return tf.func.experimental.capture(lambda: x, by_ref=False)
+
+    tf_f = tf.function(f)
+    x = type_before(val_before)
+    self.assertEqual(tf_f(), val_before)
+    x = type_after(val_after)
+    self.assertEqual(tf_f(), val_before)
+
+
+if __name__ == "__main__":
+  unittest.main()

tensorflow/core/function/integration_test/side_inputs_test.py

@@ -0,0 +1,187 @@
+# 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 side inputs in tf.function."""
+
+import unittest
+from absl.testing import parameterized
+
+import tensorflow as tf
+
+
+# TODO(panzf): Enable tests after side inputs are supported
+class SideInputsTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      (1, tf.constant, 2, tf.constant),
+      (1.0, tf.constant, 2.0, tf.constant),
+      (1, int, 2, int),
+      (1.0, float, 2.0, float),
+      (1, int, 2, tf.constant),
+      (1, tf.constant, 2, int))
+  @unittest.skip("Feature not implemented")
+  def test_direct_capture(self, val_before, type_before, val_after, type_after):
+    def f():
+      return x + 1
+
+    tf_f = tf.function(f)
+    x = type_before(val_before)
+    self.assertEqual(f(), tf_f())
+    x = type_after(val_after)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  def test_direct_capture_mutation(self):
+    def f():
+      return glob[-1] + tf.constant(0)
+
+    tf_f = tf.function(f)
+    glob = [tf.constant(1), tf.constant(2)]
+
+    self.assertEqual(f(), tf_f())
+    glob.append(tf.constant(3))
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(
+      tf.constant,
+      int)
+  def test_dict_capture_mutation_with_tensor_and_non_tensor(self, capture_type):
+    def f():
+      return d["val"]
+
+    tf_f = tf.function(f)
+    d = {"int": 1, "tensor": tf.constant(2), "val": capture_type(3)}
+    self.assertEqual(f(), tf_f())
+    d["val"] = capture_type(4)
+    self.assertEqual(f(), tf_f())
+
+  @unittest.skip("Feature not implemented")
+  @parameterized.parameters(tf.constant, int)
+  def test_capture_with_duplicate_usage(self, capture_type):
+    def f():
+      return x + x  # should capture x just once.
+
+    tf_f = tf.function(f)
+
+    x = capture_type(1)
+    self.assertEqual(f(), tf_f())
+    self.assertLen(tf_f.get_concrete_function().graph.inputs, 1)
+
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+    self.assertLen(tf_f.get_concrete_function().graph.inputs, 1)
+
+  @unittest.skip("Feature not implemented")
+  def test_local_capture(self):
+    def f():
+      x = tf.constant(0)
+      def g():
+        return x
+      return g()
+
+    tf_f = tf.function(f)
+    x = tf.constant(100)  # pylint: disable=unused-variable
+    self.assertEqual(f(), tf_f())
+    x = tf.constant(200)
+    self.assertEqual(f(), tf_f())
+
+  @parameterized.parameters(
+      tf.constant,
+      int)
+  @unittest.skip("Feature not implemented")
+  def test_capture_by_nested_function(self, capture_type):
+    def f():
+      def g():
+        return x
+      return g()
+
+    tf_f = tf.function(f)
+
+    x = capture_type(1)
+    self.assertEqual(f(), tf_f())
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+
+  @parameterized.parameters(tf.constant, int)
+  @unittest.skip("Feature not implemented")
+  def test_outer_capture_with_function_call(self, capture_type):
+    def g():
+      return x
+
+    @tf.function
+    def f():
+      return g()
+    tf_f = tf.function(f)
+
+    x = capture_type(1)
+    self.assertEqual(f(), tf_f())
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+
+  @parameterized.parameters(tf.constant, int)
+  @unittest.skip("Feature not implemented")
+  def test_outer_capture_with_nested_function_call(self, capture_type):
+    def g_factory():
+      def g():
+        return x
+      return g()
+
+    def f():
+      h = g_factory()
+      return h()
+    tf_f = tf.function(f)
+
+    x = capture_type(1)
+    self.assertEqual(f(), tf_f())
+    x = capture_type(2)
+    self.assertEqual(f(), tf_f())
+
+  @parameterized.parameters(tf.constant, int)
+  @unittest.skip("Feature not implemented")
+  def test_capture_within_function_argument(self, capture_type):
+    def g():
+      return x
+
+    def f(h):
+      return h()
+    tf_f = tf.function(f)
+
+    x = capture_type(1)
+    self.assertEqual(f(g), tf_f(g))
+    x = capture_type(2)
+    self.assertEqual(f(g), tf_f(g))
+
+  @parameterized.parameters(
+      tf.constant,
+      int)
+  @unittest.skip("Feature not implemented")
+  def test_nested_tf_function_with_capture(self, capture_type):
+    @tf.function
+    def tf_f():
+      @tf.function
+      def tf_g():
+        return x
+      return tf_g()
+
+    x = capture_type(0)
+    self.assertEqual(tf_f(), tf.constant(0))
+    x = capture_type(1)
+    self.assertEqual(tf_f(), tf.constant(0))
+    # Test the outer function doesn't have any captures
+    self.assertLen(tf_f.get_concrete_function().graph.capture, 1)
+
+
+if __name__ == "__main__":
+  unittest.main()

tensorflow/core/function/polymorphism/BUILD

@@ -0,0 +1,119 @@
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "type_dispatch",
+    srcs = [
+        "type_dispatch.py",
+    ],
+    visibility = ["//visibility:private"],
+    deps = [
+        ":function_type",
+    ],
+)
+
+py_test(
+    name = "type_dispatch_test",
+    srcs = ["type_dispatch_test.py"],
+    strict_deps = True,
+    deps = [
+        ":function_type",
+        ":type_dispatch",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/types:trace",
+    ],
+)
+
+pytype_strict_library(
+    name = "function_cache",
+    srcs = ["function_cache.py"],
+    visibility = [
+        "//tensorflow/python/eager:__subpackages__",
+    ],
+    deps = [
+        ":function_type",
+        "//tensorflow/core/function/polymorphism:type_dispatch",
+        "//tensorflow/core/function/trace_type",
+    ],
+)
+
+py_test(
+    name = "function_cache_test",
+    size = "medium",
+    srcs = ["function_cache_test.py"],
+    strict_deps = True,
+    # TODO(b/213375459): For continuous benchmark monitoring
+    visibility = ["//learning/brain/contrib/eager/python/examples:__pkg__"],
+    deps = [
+        ":function_cache",
+        #internal proto upb dep
+        "//tensorflow/core/function/polymorphism:function_type",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/types:trace",
+    ],
+)
+
+pytype_strict_library(
+    name = "function_type",
+    srcs = [
+        "function_type.py",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//tensorflow/core/function/polymorphism:function_type_proto_py",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/core/function/trace_type:serialization",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/types:trace",
+        "@absl_py//absl/logging",
+    ],
+)
+
+tf_proto_library(
+    name = "function_type_proto",
+    srcs = [
+        "function_type.proto",
+    ],
+    protodeps = [
+        "//tensorflow/core/function/trace_type:serialization_proto",
+    ],
+    visibility = ["//visibility:private"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "function_type_py_pb2",
+#     visibility = ["//visibility:private"],
+#     deps = [":function_type_proto"],
+# )
+# copybara:uncomment_end
+
+py_test(
+    name = "function_type_test",
+    srcs = ["function_type_test.py"],
+    strict_deps = True,
+    deps = [
+        ":function_type",
+        #internal proto upb dep
+        "//tensorflow/core/function/polymorphism:function_type_proto_py",
+        "//tensorflow/core/function/trace_type",
+        "//tensorflow/core/function/trace_type:serialization",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/types:trace",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/core/function/polymorphism/function_cache.py

@@ -0,0 +1,103 @@
+# 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.
+# ==============================================================================
+"""Cache to manage functions based on their FunctionType."""
+
+import collections
+from typing import Any, NamedTuple, Optional
+
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.core.function.polymorphism import type_dispatch
+
+
+class FunctionContext(NamedTuple):
+  """Contains information regarding tf.function execution context."""
+  context: Any = None
+  scope_type: Any = None
+
+
+class FunctionCache:
+  """A container for managing functions."""
+
+  __slots__ = ["_primary", "_dispatch_dict", "_garbage_collectors"]
+
+  def __init__(self):
+    # Maps (FunctionContext, FunctionType) to a function.
+    self._primary = collections.OrderedDict()
+
+    # Maps FunctionContext to a TypeDispatchTable containing FunctionTypes of
+    # that particular context.
+    self._dispatch_dict = {}
+
+  def lookup(self, function_type: function_type_lib.FunctionType,
+             context: Optional[FunctionContext] = None) -> Optional[Any]:
+    """Looks up a function based on the context and type."""
+    context = context or FunctionContext()
+    if context in self._dispatch_dict:
+      dispatch_type = self._dispatch_dict[context].dispatch(function_type)
+      if dispatch_type:
+        return self._primary[(context, dispatch_type)]
+
+    return None
+
+  def delete(self, function_type: function_type_lib.FunctionType,
+             context: Optional[FunctionContext] = None,
+             ) -> bool:
+    """Deletes a function given the context and type."""
+    context = context or FunctionContext()
+    if (context, function_type) not in self._primary:
+      return False
+
+    del self._primary[(context, function_type)]
+    self._dispatch_dict[context].delete(function_type)
+
+    return True
+
+  def add(self, fn: Any, context: Optional[FunctionContext] = None) -> None:
+    """Adds a new function using its function_type.
+
+    Args:
+      fn: The function to be added to the cache.
+      context: A FunctionContext representing the current context.
+    """
+    context = context or FunctionContext()
+    self._primary[(context, fn.function_type)] = fn
+    if context not in self._dispatch_dict:
+      self._dispatch_dict[context] = type_dispatch.TypeDispatchTable()
+
+    self._dispatch_dict[context].add_target(fn.function_type)
+
+  def generalize(
+      self, context: FunctionContext,
+      function_type: function_type_lib.FunctionType
+  ) -> function_type_lib.FunctionType:
+    """Try to generalize a FunctionType within a FunctionContext."""
+    if context in self._dispatch_dict:
+      return self._dispatch_dict[context].try_generalizing_function_type(
+          function_type)
+    else:
+      return function_type
+
+  # TODO(b/205971333): Remove this function.
+  def clear(self):
+    """Removes all functions from the cache."""
+    self._primary.clear()
+    self._dispatch_dict.clear()
+
+  def values(self):
+    """Returns a list of all functions held by this cache."""
+    return list(self._primary.values())
+
+  def __len__(self):
+    return len(self._primary)

tensorflow/core/function/polymorphism/function_cache_test.py

@@ -0,0 +1,483 @@
+# 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 function_cache."""
+
+import dataclasses
+import itertools
+import timeit
+from typing import Any, Optional
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.polymorphism import function_cache
+from tensorflow.core.function.polymorphism import function_type
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+from tensorflow.python.types import trace
+
+
+class MockGenericType(trace.TraceType):
+
+  def __init__(self, obj):
+    self._object = obj
+
+  def is_subtype_of(self, other):
+    return self == other
+
+  def most_specific_common_supertype(self, others):
+    return None
+
+  def placeholder_value(self, placeholder_context=None):
+    raise NotImplementedError
+
+  def __eq__(self, other):
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    return isinstance(other, MockGenericType) and self._object == other._object
+
+  def __hash__(self):
+    return hash(self._object)
+
+
+class MockIntGenericType(MockGenericType):
+
+  def most_specific_common_supertype(self, others):
+    if all([self._object == other._object for other in others]):
+      return MockIntGenericType(self._object)
+    else:
+      return None
+
+
+class MockSubtypeOf2(MockGenericType):
+
+  def is_subtype_of(self, other):
+    if not isinstance(other, MockGenericType):
+      return False
+
+    return other._object == 2
+
+
+class MockSupertypes2With3(MockGenericType):
+
+  def most_specific_common_supertype(self, others):
+    if self._object == 2 and isinstance(others[0]._object, int):
+      return MockSupertypes2With3(3)
+    else:
+      return None
+
+
+class MockShape(trace.TraceType):
+
+  def __init__(self, *shape: Optional[int]):
+    self.shape = shape
+
+  def is_subtype_of(self, other: "MockShape") -> bool:
+    if len(self.shape) != len(other.shape):
+      return False
+
+    if any(o is not None and s != o for s, o in zip(self.shape, other.shape)):
+      return False
+
+    return True
+
+  def most_specific_common_supertype(self, _):
+    raise NotImplementedError
+
+  def __str__(self):
+    return str(self.shape)
+
+  def __repr__(self):
+    return str(self)
+
+  def __hash__(self) -> int:
+    return hash(self.shape)
+
+  def __eq__(self, other: "MockShape") -> bool:
+    return self.shape == other.shape
+
+  def placeholder_value(self, placeholder_context):
+    raise NotImplementedError
+
+
+def make_single_param_type(type_constraint):
+  return function_type.FunctionType(
+      [
+          function_type.Parameter(
+              "x",
+              function_type.Parameter.POSITIONAL_ONLY,
+              False,
+              type_constraint,
+          )
+      ]
+  )
+
+
+@dataclasses.dataclass(frozen=True)
+class MockFunction:
+  function_type: Any
+  test_string: str
+
+
+def make_type(value):
+  typing_context = trace_type.InternalTracingContext()
+  value_type = trace_type.from_value(value, typing_context)
+  f_type = make_single_param_type(value_type)
+  return f_type
+
+
+class FunctionCacheTest(test.TestCase):
+
+  def testConcreteFunctionDictRetainsInsertedKeys(self):
+    cache = function_cache.FunctionCache()
+
+    f_type_1 = make_type(1)
+    self.assertIsNone(cache.lookup(f_type_1))
+
+    f_type_2 = make_type(2)
+    f_type_3 = make_type(3)
+
+    cache.add(MockFunction(f_type_1, "test_1"))
+    cache.add(MockFunction(f_type_2, "test_2"))
+
+    self.assertEqual(cache.lookup(f_type_1).test_string, "test_1")
+    self.assertEqual(cache.lookup(f_type_2).test_string, "test_2")
+    self.assertIsNone(cache.lookup(f_type_3))
+
+  def testClearRemovesAllConcreteFunctions(self):
+    cache = function_cache.FunctionCache()
+
+    f_type_1 = make_type(1)
+    f_type_2 = make_type(2)
+    f_type_3 = make_type(3)
+
+    cache.add(MockFunction(f_type_1, "test_1"))
+    cache.add(MockFunction(f_type_2, "test_2"))
+
+    self.assertEqual(cache.lookup(f_type_1).test_string, "test_1")
+    self.assertEqual(cache.lookup(f_type_2).test_string, "test_2")
+    self.assertIsNone(cache.lookup(f_type_3))
+
+    cache.clear()
+
+    self.assertIsNone(cache.lookup(f_type_1))
+    self.assertIsNone(cache.lookup(f_type_2))
+    self.assertIsNone(cache.lookup(f_type_3))
+
+  def testDeleteRemovesConcreteFunctions(self):
+    cache = function_cache.FunctionCache()
+    f_type_1 = make_type(1)
+    cache.add(MockFunction(f_type_1, "test_1"))
+    self.assertEqual(cache.lookup(f_type_1).test_string, "test_1")
+    cache.delete(f_type_1)
+    self.assertIsNone(cache.lookup(f_type_1))
+
+    f_type_2 = make_single_param_type(MockSubtypeOf2(2))
+    cache.add(MockFunction(f_type_2, "test_2"))
+    self.assertEqual(cache.lookup(f_type_2).test_string, "test_2")
+
+    f_type_3 = make_single_param_type(MockSubtypeOf2(3))
+    self.assertEqual(cache.lookup(f_type_3).test_string, "test_2")
+
+    cache.delete(f_type_2)
+    self.assertIsNone(cache.lookup(f_type_2))
+    self.assertIsNone(cache.lookup(f_type_3))
+
+  def testMostSpecificFunctionCacheKeyIsLookedUp(self):
+    ctx = function_cache.FunctionContext(0)
+    cache = function_cache.FunctionCache()
+    cache.add(
+        MockFunction(make_single_param_type(MockShape(1, 2, None)), "a"), ctx
+    )
+    cache.add(
+        MockFunction(make_single_param_type(MockShape(1, 2, 3)), "b"), ctx
+    )
+
+    self.assertEqual(
+        cache.lookup(
+            make_single_param_type(MockShape(1, 2, 3)), ctx
+        ).test_string,
+        "b",
+    )
+
+  def testFirstMostSpecificFunctionCacheKeyIsLookedUp(self):
+    ctx = function_cache.FunctionContext(0)
+    cache = function_cache.FunctionCache()
+    cache.add(
+        MockFunction(make_single_param_type(MockShape(1, 2, None)), "a"), ctx
+    )
+    cache.add(
+        MockFunction(
+            make_single_param_type(MockShape(1, None, 3)),
+            "b",
+        ),
+        ctx,
+    )
+
+    self.assertEqual(
+        cache.lookup(
+            make_single_param_type(MockShape(1, 2, 3)), ctx
+        ).test_string,
+        "a",
+    )
+
+  def testMostSpecificFunctionCacheKeyIsOrderAgnostic(self):
+    ctx = function_cache.FunctionContext(0)
+    keys = [
+        (MockFunction(make_single_param_type(MockShape(1, 1, 1)), "a"), ctx),
+        (MockFunction(make_single_param_type(MockShape(1, None, 1)), "b"), ctx),
+        (
+            MockFunction(make_single_param_type(MockShape(None, None, 1)), "c"),
+            ctx,
+        ),
+        (
+            MockFunction(
+                make_single_param_type(MockShape(None, None, None)), "d"
+            ),
+            ctx,
+        ),
+    ]
+
+    for permutation in itertools.permutations(keys):
+      cache = function_cache.FunctionCache()
+      cache.add(
+          permutation[0][0],
+          permutation[0][1],
+      )
+      cache.add(
+          permutation[1][0],
+          permutation[1][1],
+      )
+      cache.add(
+          permutation[2][0],
+          permutation[2][1],
+      )
+      cache.add(
+          permutation[3][0],
+          permutation[3][1],
+      )
+
+      self.assertEqual(
+          cache.lookup(
+              make_single_param_type(MockShape(1, 1, 1)), ctx
+          ).test_string,
+          "a",
+      )
+      self.assertEqual(
+          cache.lookup(
+              make_single_param_type(MockShape(1, 2, 1)), ctx
+          ).test_string,
+          "b",
+      )
+      self.assertEqual(
+          cache.lookup(
+              make_single_param_type(MockShape(2, 2, 1)), ctx
+          ).test_string,
+          "c",
+      )
+      self.assertEqual(
+          cache.lookup(
+              make_single_param_type(MockShape(2, 2, 2)), ctx
+          ).test_string,
+          "d",
+      )
+
+
+class FunctionCacheBenchmark(test.Benchmark):
+
+  def benchmarkCacheHit50thKeyMiss(self):
+    # If there are 50 keys and we get a new key that the cache has no concrete
+    # functions for.
+
+    cache = function_cache.FunctionCache()
+    args_per_call = 5
+    num_total_checks = 50
+
+    keys = []
+    for i in range(num_total_checks):
+      args = []
+      for j in range(args_per_call):
+        args.append(array_ops.zeros([i, j]))
+      keys.append(make_type(args))
+
+    for key in keys[:-1]:
+      cache.add(MockFunction(key, "testing"))
+
+    iterations = 10000
+    subtyping_time = timeit.timeit(
+        lambda: cache.lookup(keys[-1]),
+        number=iterations,
+    )
+    equality_time = timeit.timeit(
+        lambda: cache.lookup(keys[-1]),
+        number=iterations,
+    )
+
+    self.report_benchmark(
+        name="cache_hit_50th_f_type_miss",
+        iters=iterations,
+        wall_time=subtyping_time + equality_time,
+        metrics=[
+            {
+                "name": "cache_hit_50th_f_type_miss_subtype_avg_ms",
+                "value": subtyping_time / iterations * 1000,
+            },
+            {
+                "name": "cache_hit_50th_f_type_miss_equality_avg_ms",
+                "value": equality_time / iterations * 1000,
+            },
+            {
+                "name": (
+                    "cache_hit_50th_f_type_miss_subtype_over_equality_ratio"
+                ),
+                "value": subtyping_time / equality_time,
+            },
+        ],
+    )
+
+  def benchmarkCacheHit50thKeyEqual(self):
+    # If there are 50 keys and we get a new key that is equal to a key that is
+    # in the cache.
+
+    cache = function_cache.FunctionCache()
+    args_per_call = 5
+    num_total_checks = 50
+
+    keys = []
+    for i in range(num_total_checks):
+      args = []
+      for j in range(args_per_call):
+        args.append(array_ops.zeros([i, j]))
+      keys.append(make_type(args))
+
+    for key in keys:
+      cache.add(MockFunction(key, "testing"))
+
+    iterations = 10000
+    subtyping_time = timeit.timeit(
+        lambda: cache.lookup(keys[-1]),
+        number=iterations,
+    )
+    equality_time = timeit.timeit(
+        lambda: cache.lookup(keys[-1]),
+        number=iterations,
+    )
+
+    self.report_benchmark(
+        name="cache_hit_50th_f_type_equal",
+        iters=iterations,
+        wall_time=subtyping_time + equality_time,
+        metrics=[
+            {
+                "name": "cache_hit_50th_f_type_equal_subtype_avg_ms",
+                "value": subtyping_time / iterations * 1000,
+            },
+            {
+                "name": "cache_hit_50th_f_type_equal_equality_avg_ms",
+                "value": equality_time / iterations * 1000,
+            },
+            {
+                "name": "cache_hit_50th_f_type_subtype_over_equality_ratio",
+                "value": subtyping_time / equality_time,
+            },
+        ],
+    )
+
+  def benchmarkCacheHit50thKeyKnownSubtype(self):
+    # If there are 50 keys and we get a key that has a subtype in cache and
+    # the cache has observed the key before (to memorize the subtype).
+
+    cache = function_cache.FunctionCache()
+    args_per_call = 5
+    num_total_checks = 50
+
+    keys = []
+    for i in range(num_total_checks - 1):
+      args = []
+      for j in range(args_per_call):
+        args.append(array_ops.zeros([i, j]))
+      keys.append(make_type(args))
+
+    for key in keys:
+      cache.add(MockFunction(key, "testing"))
+    cache.add(
+        MockFunction(make_single_param_type(MockSubtypeOf2(2)), "testing"),
+    )
+    cache.lookup(make_single_param_type(MockSubtypeOf2(3)))
+
+    iterations = 10000
+    lookup_key = make_single_param_type(MockSubtypeOf2(2))
+    subtyping_time = timeit.timeit(
+        lambda: cache.lookup(lookup_key), number=iterations
+    )
+
+    self.report_benchmark(
+        name="cache_hit_50th_f_type_known_subtype",
+        iters=iterations,
+        wall_time=subtyping_time,
+        metrics=[{
+            "name": "cache_hit_50th_f_type_known_subtype_avg_ms",
+            "value": subtyping_time / iterations * 1000,
+        }],
+    )
+
+  def benchmarkCacheHit50thKeyUnknownSubtype(self):
+    # If there are 50 keys and we get a key that has a subtype in cache but
+    # the cache has never observed the key before (no memory for the subtype).
+
+    cache = function_cache.FunctionCache()
+    args_per_call = 5
+    num_total_checks = 50
+
+    keys = []
+    for i in range(num_total_checks - 1):
+      args = []
+      for j in range(args_per_call):
+        args.append(array_ops.zeros([i, j]))
+      keys.append(make_type(args))
+
+    def setup():
+      cache.clear()
+      for key in keys:
+        cache.add(
+            MockFunction(key, "testing"),
+        )
+      cache.add(
+          MockFunction(make_single_param_type(MockSubtypeOf2(3)), "testing"),
+      )
+
+    iterations = 10000
+    lookup_key = make_single_param_type(MockSubtypeOf2(2))
+    subtyping_time = sum(
+        timeit.repeat(
+            stmt=lambda: cache.lookup(lookup_key),
+            setup=setup,
+            repeat=iterations,
+            number=1,
+        )
+    )
+
+    self.report_benchmark(
+        name="cache_hit_50th_f_type_unknown_subtype",
+        iters=iterations,
+        wall_time=subtyping_time,
+        metrics=[{
+            "name": "cache_hit_50th_f_type_unknown_subtype_avg_ms",
+            "value": subtyping_time / iterations * 1000,
+        }],
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/core/function/polymorphism/function_type.proto

@@ -0,0 +1,51 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto2";
+
+package tensorflow.core.function.polymorphism.function_type;
+
+import "tensorflow/core/function/trace_type/serialization.proto";
+
+// Represents a serialized Parameter type.
+message Parameter {
+  enum Kind {
+    UNDEFINED = 0;
+    POSITIONAL_ONLY = 1;
+    POSITIONAL_OR_KEYWORD = 2;
+    VAR_POSITIONAL = 3;
+    KEYWORD_ONLY = 4;
+    VAR_KEYWORD = 5;
+  }
+  optional string name = 1;
+  optional Kind kind = 2;
+  optional bool is_optional = 3;
+  optional tensorflow.core.function.trace_type.serialization.SerializedTraceType
+      type_constraint = 4;
+}
+
+// Represents a serialized Capture type.
+message Capture {
+  optional string name = 1;
+  optional tensorflow.core.function.trace_type.serialization.SerializedTraceType
+      type_constraint = 2;
+}
+
+// Represents a serialized FunctionType.
+message FunctionType {
+  repeated Parameter parameters = 1;
+  repeated Capture captures = 2;
+  // TODO(fmuham): Add support for return type.
+}

tensorflow/core/function/polymorphism/function_type.py

@@ -0,0 +1,733 @@
+# 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.
+# ==============================================================================
+"""Represents the types of TF functions."""
+
+import collections
+import inspect
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple
+
+from absl import logging
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.polymorphism import function_type_pb2
+from tensorflow.core.function.trace_type import serialization
+from tensorflow.python.types import core
+from tensorflow.python.types import trace
+
+
+# Represents a defined parameter default value that is saved alongside the
+# function's captures.
+class CapturedDefaultValue:
+  def __repr__(self):
+    return "<captured_default_value>"
+
+  def __str__(self):
+    return "<captured_default_value>"
+
+CAPTURED_DEFAULT_VALUE = CapturedDefaultValue()
+
+PROTO_TO_PY_ENUM = {
+    function_type_pb2.Parameter.Kind.POSITIONAL_ONLY:
+        inspect.Parameter.POSITIONAL_ONLY,
+    function_type_pb2.Parameter.Kind.POSITIONAL_OR_KEYWORD:
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    function_type_pb2.Parameter.Kind.VAR_POSITIONAL:
+        inspect.Parameter.VAR_POSITIONAL,
+    function_type_pb2.Parameter.Kind.KEYWORD_ONLY:
+        inspect.Parameter.KEYWORD_ONLY,
+    function_type_pb2.Parameter.Kind.VAR_KEYWORD:
+        inspect.Parameter.VAR_KEYWORD,
+}
+
+PY_TO_PROTO_ENUM = {v: k for k, v in PROTO_TO_PY_ENUM.items()}
+
+
+class Parameter(inspect.Parameter):
+  """Represents a parameter to a function."""
+
+  def __init__(self, name: str, kind: Any, optional: bool,
+               type_constraint: Optional[trace.TraceType]):
+    if optional and kind not in [
+        self.POSITIONAL_ONLY, self.KEYWORD_ONLY, self.POSITIONAL_OR_KEYWORD
+    ]:
+      raise ValueError(
+          "Parameter " + name +
+          " is optional and its kind must be one of {POSITIONAL_ONLY, " +
+          "KEYWORD_ONLY, POSITIONAL_OR_KEYWORD}. Got: " + str(kind))
+
+    if type_constraint and kind in [self.VAR_POSITIONAL, self.VAR_KEYWORD]:
+      raise TypeError("Variable args/kwargs can not have type constraints.")
+
+    if not isinstance(type_constraint, (trace.TraceType, type(None))):
+      raise TypeError(
+          "Type constraints can only be an instance of a TraceType but got " +
+          "type_constraint=" + str(type_constraint) + " for Parameter " + name)
+
+    super().__init__(
+        name,
+        kind,
+        default=CAPTURED_DEFAULT_VALUE if optional else self.empty,
+        annotation=type_constraint
+        if type_constraint is not None else self.empty)
+
+  @classmethod
+  def from_proto(cls, proto: Any) -> "Parameter":
+    """Generate a Parameter from the proto representation."""
+    deserialized_type_constraint = serialization.deserialize(
+        proto.type_constraint) if proto.HasField("type_constraint") else None
+    return Parameter(proto.name, PROTO_TO_PY_ENUM[proto.kind],
+                     proto.is_optional, deserialized_type_constraint)
+
+  def to_proto(self) -> function_type_pb2.Parameter:
+    """Generate a proto representation of the Parameter."""
+    serialized_type_constraint = serialization.serialize(
+        self.type_constraint) if self.type_constraint else None
+    return function_type_pb2.Parameter(
+        name=self.name,
+        kind=PY_TO_PROTO_ENUM[self.kind],
+        is_optional=self.optional,
+        type_constraint=serialized_type_constraint)
+
+  @property
+  def optional(self) -> bool:
+    """If this parameter might not be supplied for a call."""
+    return self.default is not self.empty
+
+  @property
+  def type_constraint(self) -> Optional[trace.TraceType]:
+    """A supertype that the parameter's type must subtype for validity."""
+    return self.annotation if self.annotation is not self.empty else None
+
+  def is_subtype_of(self, other: "Parameter") -> bool:
+    """Returns True if self is a supertype of other Parameter."""
+    if not self.type_constraint or not other.type_constraint:
+      raise TypeError(
+          "Can not determine relationship between partially specified types.")
+
+    if ((self.name, self.kind, self.optional) !=
+        (other.name, other.kind, other.optional)):
+      return False
+
+    return self.type_constraint.is_subtype_of(other.type_constraint)
+
+  def most_specific_common_supertype(
+      self, others: Sequence["Parameter"]) -> Optional["Parameter"]:
+    """Returns a common supertype (if exists)."""
+    if not self.type_constraint or any(
+        not other.type_constraint for other in others):
+      raise TypeError(
+          "Can not determine relationship between partially specified types.")
+
+    for other in others:
+      if ((self.name, self.kind, self.optional) !=
+          (other.name, other.kind, other.optional)):
+        return None
+
+    supertyped_constraint = self.type_constraint.most_specific_common_supertype(
+        [other.type_constraint for other in others])
+    if supertyped_constraint:
+      return Parameter(self.name, self.kind, self.optional,
+                       supertyped_constraint)
+    else:
+      return None
+
+  def __eq__(self, other: Any) -> bool:
+    if not isinstance(other, Parameter):
+      return NotImplemented
+
+    return ((self.name, self.kind, self.optional,
+             self.type_constraint) == (other.name, other.kind, other.optional,
+                                       other.type_constraint))
+
+  def __hash__(self):
+    return hash((self.name, self.kind, self.optional, self.type_constraint))
+
+  def __repr__(self):
+    return ("Parameter(name=" + self.name + ", kind=" + str(self.kind) +
+            ", optional=" + repr(self.optional) + ", type_constraint=" +
+            repr(self.type_constraint) + ")")
+
+  def __reduce__(self):
+    return (self.__class__, (self.name, self.kind, self.optional,
+                             self.type_constraint))
+
+
+class FunctionType(core.FunctionType):
+  """Represents the type of a TensorFlow function.
+
+  FunctionType is the canonical way to represent the input/output contract of
+  all kinds of functions within the tf.function domain, including:
+    - Polymorphic Function
+    - Concrete Function
+    - Atomic Function
+
+  It provides consistent, centralized and layered logic for:
+    - Canonicalization of Python input arguments
+    - Type-based dispatch to monomorphic functions
+    - Packing/unpacking structured python values to Tensors
+    - Generation of structured placeholder values for tracing
+
+  Additionaly, it also provides:
+    - Lossless serialization
+    - Native integration with Python function signature representation
+    - Seamless migration from older representation formats
+  """
+
+  def __init__(self,
+               parameters: Sequence[inspect.Parameter],
+               captures: Optional[collections.OrderedDict] = None,
+               **kwargs):
+    super().__init__(parameters, **kwargs)
+    self._captures = captures if captures else collections.OrderedDict()
+
+  @property
+  def parameters(self) -> Mapping[str, Any]:
+    """Returns an ordered mapping of parameter name to specification."""
+    return super().parameters
+
+  @property
+  def captures(self) -> collections.OrderedDict:
+    """Returns an ordered mapping of capture id to type."""
+    return self._captures
+
+  @property
+  def output(self) -> Optional[trace.TraceType]:
+    """Return the output TraceType if specified."""
+    return (
+        self.return_annotation
+        if self.return_annotation is not self.empty
+        else None
+    )
+
+  @classmethod
+  def from_callable(cls,
+                    obj: Callable[..., Any],
+                    *,
+                    follow_wrapped: bool = True) -> "FunctionType":
+    """Generate FunctionType from a python Callable."""
+    signature = super().from_callable(obj, follow_wrapped=follow_wrapped)
+    # TODO(fmuham): Support TraceType-based annotations.
+    parameters = [
+        Parameter(p.name, p.kind, p.default is not p.empty, None)
+        for p in signature.parameters.values()
+    ]
+
+    return FunctionType(parameters)
+
+  @classmethod
+  def get_default_values(cls,
+                         obj: Callable[..., Any],
+                         *,
+                         follow_wrapped: bool = True) -> Dict[str, Any]:
+    """Inspects and returns a dictionary of default values."""
+    signature = super().from_callable(obj, follow_wrapped=follow_wrapped)
+    default_values = {}
+    for p in signature.parameters.values():
+      if p.default is not p.empty:
+        default_values[p.name] = p.default
+    return default_values
+
+  @classmethod
+  def from_proto(cls, proto: Any) -> "FunctionType":
+    """Generate a FunctionType from the proto representation."""
+    return FunctionType([Parameter.from_proto(p) for p in proto.parameters],
+                        collections.OrderedDict([
+                            (c.name,
+                             serialization.deserialize(c.type_constraint))
+                            for c in proto.captures
+                        ]))
+
+  def to_proto(self) -> Any:
+    """Generate a proto representation from the FunctionType."""
+    return function_type_pb2.FunctionType(
+        parameters=[p.to_proto() for p in self.parameters.values()],
+        captures=[
+            function_type_pb2.Capture(
+                name=n, type_constraint=serialization.serialize(t))
+            for n, t in self.captures.items()
+        ])
+
+  def bind_with_defaults(self, args, kwargs, default_values):
+    """Returns BoundArguments with default values filled in."""
+    bound_arguments = self.bind(*args, **kwargs)
+    bound_arguments.apply_defaults()
+
+    with_default_args = collections.OrderedDict()
+    for name, value in bound_arguments.arguments.items():
+      if value is CAPTURED_DEFAULT_VALUE:
+        with_default_args[name] = default_values[name]
+      else:
+        with_default_args[name] = value
+
+    for arg_name in with_default_args:
+      constraint = self.parameters[arg_name].type_constraint
+      if constraint:
+        with_default_args[arg_name] = constraint.cast(
+            with_default_args[arg_name],
+            trace_type.InternalCastContext(allow_specs=True),
+        )
+    bound_arguments = inspect.BoundArguments(self, with_default_args)
+    return bound_arguments
+
+  def is_supertype_of(self, other: "FunctionType") -> bool:
+    """Returns True if self is a supertype of other FunctionType."""
+    if len(self.parameters) != len(other.parameters):
+      return False
+
+    for self_param, other_param in zip(self.parameters.values(),
+                                       other.parameters.values()):
+      # Functions are contravariant on their parameter types.
+      if not self_param.is_subtype_of(other_param):
+        return False
+
+    # Other must have all capture names of self.
+    if not all(name in other.captures for name in self.captures):
+      return False
+
+    # Functions are contravariant upon the capture types.
+    return all(capture_type.is_subtype_of(other.captures[name])
+               for name, capture_type in self.captures.items())
+
+  def most_specific_common_subtype(
+      self, others: Sequence["FunctionType"]) -> Optional["FunctionType"]:
+    """Returns a common subtype (if exists)."""
+    subtyped_parameters = []
+
+    for i, parameter in enumerate(self.parameters.values()):
+      # Functions are contravariant on their parameter types.
+      subtyped_parameter = parameter.most_specific_common_supertype(
+          [list(other.parameters.values())[i] for other in others])
+      if subtyped_parameter is None:
+        return None
+      subtyped_parameters.append(subtyped_parameter)
+
+    if not all(subtyped_parameters):
+      return None
+
+    # Common subtype has superset of all captures.
+    capture_names = set(self.captures.keys())
+    for other in others:
+      capture_names = capture_names.union(other.captures.keys())
+
+    subtyped_captures = collections.OrderedDict()
+    for name in capture_names:
+      containing = [t for t in [self, *others] if name in t.captures]
+      # Pick the first type that has the capture as the base.
+      base = containing[0]
+      relevant_others = containing[1:]
+
+      # Functions are contravariant upon the capture types.
+      common_type = base.captures[name].most_specific_common_supertype(
+          [other.captures[name] for other in relevant_others]
+      )
+      if common_type is None:
+        return None
+      else:
+        subtyped_captures[name] = common_type
+
+    return FunctionType(subtyped_parameters, subtyped_captures)
+
+  def placeholder_arguments(
+      self, placeholder_context: trace.PlaceholderContext
+  ) -> inspect.BoundArguments:
+    """Returns BoundArguments of values that can be used for tracing."""
+    arguments = collections.OrderedDict()
+    for parameter in self.parameters.values():
+      if parameter.kind in {Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD}:
+        raise ValueError("Can not generate placeholder values for "
+                         "variable length function type.")
+
+      if not parameter.type_constraint:
+        raise ValueError("Can not generate placeholder value for "
+                         "partially defined function type.")
+      placeholder_context.update_naming_scope(parameter.name)
+      arguments[parameter.name] = parameter.type_constraint.placeholder_value(
+          placeholder_context)
+
+    return inspect.BoundArguments(self, arguments)
+
+  @property
+  def _sorted_parameters(self) -> List[trace.TraceType]:
+    if not hasattr(self, "_cached_sorted_parameters"):
+      # Sort keyword-only parameters by name.
+      sorted_parameters = []
+      kwonly_parameters = []
+      for p in self.parameters.values():
+        if p.kind is Parameter.KEYWORD_ONLY:
+          kwonly_parameters.append(p)
+        else:
+          sorted_parameters.append(p)
+      sorted_parameters = sorted_parameters + sorted(
+          kwonly_parameters, key=lambda p: p.name
+      )
+      self._cached_sorted_parameters = sorted_parameters
+    return self._cached_sorted_parameters
+
+  @property
+  def flat_inputs(self) -> List[trace.TraceType]:
+    """Flat tensor inputs accepted by this FunctionType."""
+    if not hasattr(self, "_cached_flat_inputs"):
+      cached_flat_inputs = []
+      for p in self._sorted_parameters:
+        cached_flat_inputs.extend(p.type_constraint.flatten())
+      self._cached_flat_inputs = cached_flat_inputs
+    return self._cached_flat_inputs
+
+  def unpack_inputs(
+      self, bound_parameters: inspect.BoundArguments
+  ) -> List[core.Tensor]:
+    """Unpacks python arguments to flat tensor inputs accepted by this type."""
+    flat = []
+    for p in self._sorted_parameters:
+      flat.extend(
+          p.type_constraint.to_tensors(bound_parameters.arguments[p.name])
+      )
+
+    dealiased_inputs = []
+    ids_used = set()
+    for tensor, input_type in zip(flat, self.flat_inputs):
+      alias_id = input_type._alias_id()  # pylint: disable=protected-access
+      if alias_id is None or alias_id not in ids_used:
+        dealiased_inputs.append(tensor)
+
+      if alias_id is not None:
+        ids_used.add(alias_id)
+
+    return dealiased_inputs
+
+  @property
+  def flat_captures(self) -> List[trace.TraceType]:
+    """Flat tensor captures needed by this FunctionType."""
+    if not hasattr(self, "_cached_flat_captures"):
+      cached_flat_captures = []
+      for t in self.captures.values():
+        cached_flat_captures.extend(t.flatten())
+      self._cached_flat_captures = cached_flat_captures
+
+    return self._cached_flat_captures
+
+  def unpack_captures(self, captures) -> List[core.Tensor]:
+    """Unpacks captures to flat tensors."""
+    flat = []
+    for v, t in zip(captures, self.captures.values()):
+      flat.extend(t.to_tensors(v))
+    if len(flat) != len(self.flat_captures):
+      raise TypeError(
+          f"Flattening captures {captures} with type {self!r} produced"
+          f" {len(flat)} tensors instead of {len(self.flat_captures)}"
+      )
+    return flat
+
+  @property
+  def flat_outputs(self) -> List[trace.TraceType]:
+    """Flat tensor outputs returned by this FunctionType."""
+    if not hasattr(self, "_cached_flat_outputs"):
+      if self.output is not None:
+        self._cached_flat_outputs = self.output.flatten()
+
+    return self._cached_flat_outputs
+
+  def pack_output(self, flat_values: Sequence[core.Tensor]) -> Any:
+    """Packs flat tensors to generate a value of the output type."""
+    if flat_values is None:
+      flat_values = []
+
+    if self.output is None:
+      raise ValueError("Can not pack outputs for undefined output type.")
+    else:
+      return self.output.from_tensors(iter(flat_values))
+
+  def __eq__(self, other: Any) -> bool:
+    if not isinstance(other, FunctionType):
+      return NotImplemented
+
+    return (self.parameters, self.captures) == (other.parameters,
+                                                other.captures)
+
+  def __hash__(self) -> int:
+    return hash((tuple(self.parameters.items()), tuple(self.captures.items())))
+
+  def __repr__(self):
+    if hasattr(self, "_cached_repr"):
+      return self._cached_repr
+
+    lines = ["Input Parameters:"]
+    for parameter in self.parameters.values():
+      lines.append(
+          f"  {parameter.name} ({parameter.kind}): {parameter.type_constraint}"
+      )
+
+    lines.append("Output Type:")
+    lines.append(f"  {self.output}")
+
+    lines.append("Captures:")
+    if self.captures:
+      for capture_id, capture_type in self.captures.items():
+        lines.append(f"  {capture_id}: {capture_type}")
+    else:
+      lines.append("  None")
+
+    self._cached_repr = "\n".join(lines)
+    return self._cached_repr
+
+
+MAX_SANITIZATION_WARNINGS = 5
+sanitization_warnings_given = 0
+
+
+# TODO(fmuham): In future, replace warning with exception.
+# TODO(fmuham): Sanitize to graph node conventions.
+def sanitize_arg_name(name: str) -> str:
+  """Sanitizes function argument names.
+
+  Matches Python symbol naming rules.
+
+  Without sanitization, names that are not legal Python parameter names can be
+  set which makes it challenging to represent callables supporting the named
+  calling capability.
+
+  Args:
+    name: The name to sanitize.
+
+  Returns:
+    A string that meets Python parameter conventions.
+  """
+  # Replace non-alphanumeric chars with '_'
+  swapped = "".join([c if c.isalnum() else "_" for c in name])
+  result = swapped if swapped[0].isalpha() else "arg_" + swapped
+
+  global sanitization_warnings_given
+  if name != result and sanitization_warnings_given < MAX_SANITIZATION_WARNINGS:
+    logging.warning(
+        "`%s` is not a valid tf.function parameter name. Sanitizing to `%s`.",
+        name, result)
+    sanitization_warnings_given += 1
+
+  return result
+
+
+# TODO(fmuham): Consider forcing kind to be always POSITIONAL_OR_KEYWORD.
+def _make_validated_mono_param(
+    name, value, kind, type_context, poly_type
+) -> Parameter:
+  """Generates and validates a parameter for Monomorphic FunctionType."""
+  mono_type = trace_type.from_value(value, type_context)
+
+  if poly_type and not mono_type.is_subtype_of(poly_type):
+    raise TypeError(f"Parameter `{name}` was expected to be of type "
+                    f"{poly_type} but is {mono_type}")
+
+  return Parameter(name, kind, False, mono_type)
+
+
+def canonicalize_to_monomorphic(
+    args: Tuple[Any, ...], kwargs: Dict[Any, Any], default_values: Dict[Any,
+                                                                        Any],
+    capture_types: collections.OrderedDict, polymorphic_type: FunctionType
+) -> Tuple[FunctionType, trace_type.InternalTracingContext]:
+  """Generates a monomorphic type out of polymorphic type for given args."""
+  poly_bound_arguments = polymorphic_type.bind(*args, **kwargs)
+
+  # Inject Default Values.
+  if default_values:
+    poly_bound_arguments.apply_defaults()
+    default_values_injected = poly_bound_arguments.arguments
+    for name, value in default_values_injected.items():
+      if value is CAPTURED_DEFAULT_VALUE:
+        default_values_injected[name] = default_values[name]
+    poly_bound_arguments = inspect.BoundArguments(
+        poly_bound_arguments.signature, default_values_injected
+    )
+
+  parameters = []
+  type_context = trace_type.InternalTracingContext()
+  has_var_positional = any(p.kind is Parameter.VAR_POSITIONAL
+                           for p in polymorphic_type.parameters.values())
+
+  for name, arg in poly_bound_arguments.arguments.items():
+    poly_parameter = polymorphic_type.parameters[name]
+    if (has_var_positional and
+        poly_parameter.kind is Parameter.POSITIONAL_OR_KEYWORD):
+      # If there is a VAR_POSITIONAL, all POSITIONAL_OR_KEYWORD become
+      # POSITIONAL_ONLY.
+      parameters.append(
+          _make_validated_mono_param(name, arg, Parameter.POSITIONAL_ONLY,
+                                     type_context,
+                                     poly_parameter.type_constraint))
+
+    elif poly_parameter.kind is Parameter.VAR_POSITIONAL:
+      # Unbundle VAR_POSITIONAL into individual POSITIONAL_ONLY args.
+      for i, value in enumerate(arg):
+        parameters.append(
+            _make_validated_mono_param(f"{poly_parameter.name}_{i}", value,
+                                       Parameter.POSITIONAL_ONLY, type_context,
+                                       poly_parameter.type_constraint))
+
+    elif poly_parameter.kind is Parameter.VAR_KEYWORD:
+      # Unbundle VAR_KEYWORD into individual KEYWORD_ONLY args.
+      for kwarg_name in sorted(arg.keys()):
+        parameters.append(
+            _make_validated_mono_param(kwarg_name, arg[kwarg_name],
+                                       Parameter.KEYWORD_ONLY, type_context,
+                                       poly_parameter.type_constraint))
+    else:
+      parameters.append(
+          _make_validated_mono_param(name, arg, poly_parameter.kind,
+                                     type_context,
+                                     poly_parameter.type_constraint))
+
+  return FunctionType(parameters, capture_types), type_context
+
+
+# TODO(fmuham): Share code with canonicalize_to_monomorphic.
+# TODO(fmuham): Lift unnecessary restrictions on input_signature validity.
+def add_type_constraints(function_type: FunctionType, input_signature: Any,
+                         default_values: Dict[str, Any]) -> FunctionType:
+  """Adds type constraints to a FunctionType based on the input_signature."""
+  context = trace_type.InternalTracingContext(is_legacy_signature=True)
+  constraints = [trace_type.from_value(c, context) for c in input_signature]
+  parameters = []
+
+  has_var_pos = any(
+      p.kind is p.VAR_POSITIONAL for p in function_type.parameters.values())
+
+  for param in function_type.parameters.values():
+    # VAR_POSITIONAL does not allow POSITIONAL_OR_KEYWORD args.
+    sanitized_kind = (
+        param.POSITIONAL_ONLY if has_var_pos and
+        param.kind is param.POSITIONAL_OR_KEYWORD else param.kind)
+
+    if param.name == "self":
+      # Type constraints do not apply on them.
+      parameters.append(Parameter("self", sanitized_kind, param.optional, None))
+
+    elif param.kind is param.VAR_KEYWORD:
+      # Disabled when input_signature is specified.
+      continue
+
+    elif param.kind is param.VAR_POSITIONAL:
+      # Convert into Positional Only args based on length of constraints.
+      for i in range(len(constraints)):
+        parameters.append(
+            Parameter(param.name + "_" + str(i), Parameter.POSITIONAL_ONLY,
+                      False, constraints.pop(0)))
+
+    elif (param.kind in [
+        param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD, param.KEYWORD_ONLY
+    ]):
+      if param.kind is param.KEYWORD_ONLY and param.name not in default_values:
+        raise TypeError(
+            "Since input_signature is defined, keyword-only parameter"
+            f" `{param.name}` must have a default value"
+        )
+
+      if constraints:
+        parameters.append(
+            Parameter(param.name, sanitized_kind, param.optional,
+                      constraints.pop(0)))
+      elif param.name in default_values:
+        type_constraint = trace_type.from_value(default_values[param.name])
+        parameters.append(
+            Parameter(param.name, sanitized_kind, param.optional,
+                      type_constraint))
+      else:
+        raise TypeError(
+            f"input_signature missing type constraint for {param.name}")
+
+  if constraints:
+    raise TypeError(
+        f"input_signature contains {len(constraints)} extra type constraints.")
+
+  return FunctionType(parameters)
+
+
+def from_structured_signature(
+    input_signature=None,
+    output_signature=None,
+    capture_types=None,
+    are_keyword_args_also_positional=False,
+) -> FunctionType:
+  """Generates a FunctionType from legacy signature representation."""
+  if input_signature is None:
+    input_signature = ((), {})
+
+  args, kwargs = input_signature
+  parameters = []
+
+  for i, arg in enumerate(args):
+    parameters.append(
+        Parameter(
+            "arg_" + str(i),
+            Parameter.POSITIONAL_ONLY,
+            False,
+            trace_type.from_value(
+                arg, trace_type.InternalTracingContext(is_legacy_signature=True)
+            ),
+        )
+    )
+
+  keyword_arg_kind = (
+      Parameter.POSITIONAL_OR_KEYWORD
+      if are_keyword_args_also_positional
+      else Parameter.KEYWORD_ONLY
+  )
+
+  for name, kwarg in kwargs.items():
+    parameters.append(
+        Parameter(
+            sanitize_arg_name(name),
+            keyword_arg_kind,
+            False,
+            trace_type.from_value(
+                kwarg,
+                trace_type.InternalTracingContext(is_legacy_signature=True),
+            ),
+        )
+    )
+
+  return_type = trace_type.from_value(
+      output_signature,
+      trace_type.InternalTracingContext(is_legacy_signature=True),
+  )
+
+  return FunctionType(
+      parameters, capture_types or {}, return_annotation=return_type
+  )
+
+
+def to_structured_signature(function_type: FunctionType) -> Tuple[Any, Any]:
+  """Returns structured input and output signatures from a FunctionType."""
+  def to_signature(x_type):
+    if x_type is None:
+      raise TypeError(
+          "Can not generate structured signature if FunctionType is not fully"
+          f" specified. Received {function_type}"
+      )
+    return x_type.placeholder_value(
+        trace_type.InternalPlaceholderContext(unnest_only=True)
+    )
+
+  args_signature = []
+  kwargs_signature = {}
+  for p in function_type.parameters.values():
+    if p.kind == Parameter.POSITIONAL_ONLY:
+      args_signature.append(to_signature(p.type_constraint))
+    else:
+      kwargs_signature[p.name] = to_signature(p.type_constraint)
+
+  input_signature = (tuple(args_signature), kwargs_signature)
+  output_signature = to_signature(function_type.output)
+
+  return input_signature, output_signature

tensorflow/core/function/polymorphism/function_type_test.py

@@ -0,0 +1,866 @@
+# 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."""
+
+import collections
+import pickle
+import platform
+
+from absl.testing import parameterized
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.polymorphism import function_type
+from tensorflow.core.function.polymorphism import function_type_pb2
+from tensorflow.core.function.trace_type import serialization
+from tensorflow.python.framework import func_graph
+from tensorflow.python.platform import test
+from tensorflow.python.types import trace
+
+
+class FunctionTypeTest(test.TestCase):
+
+  def test_required_only(self):
+
+    def foo(x, y, z):  # pylint: disable=unused-argument
+      pass
+
+    constraint = function_type.FunctionType.from_callable(foo)
+    self.assertEqual(
+        constraint,
+        function_type.FunctionType(
+            (function_type.Parameter(
+                "x", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                None),
+             function_type.Parameter(
+                 "y", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                 None),
+             function_type.Parameter(
+                 "z", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                 None))))
+    self.assertEqual(function_type.FunctionType.get_default_values(foo), {})
+
+  def test_optional_only(self):
+
+    def foo(x=1, y=2, z=3):  # pylint: disable=unused-argument
+      pass
+
+    constraint = function_type.FunctionType.from_callable(foo)
+    self.assertEqual(
+        constraint,
+        function_type.FunctionType(
+            (function_type.Parameter(
+                "x", function_type.Parameter.POSITIONAL_OR_KEYWORD, True, None),
+             function_type.Parameter(
+                 "y", function_type.Parameter.POSITIONAL_OR_KEYWORD, True,
+                 None),
+             function_type.Parameter(
+                 "z", function_type.Parameter.POSITIONAL_OR_KEYWORD, True,
+                 None))))
+    self.assertEqual(
+        function_type.FunctionType.get_default_values(foo), {
+            "x": 1,
+            "y": 2,
+            "z": 3
+        })
+
+  def test_required_and_optional(self):
+
+    def foo(x, y, z=3):  # pylint: disable=unused-argument
+      pass
+
+    constraint = function_type.FunctionType.from_callable(foo)
+    self.assertEqual(
+        constraint,
+        function_type.FunctionType(
+            (function_type.Parameter(
+                "x", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                None),
+             function_type.Parameter(
+                 "y", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                 None),
+             function_type.Parameter(
+                 "z", function_type.Parameter.POSITIONAL_OR_KEYWORD, True,
+                 None))))
+    self.assertEqual(
+        function_type.FunctionType.get_default_values(foo), {"z": 3})
+
+  def test_method_bound(self):
+
+    class MyClass:
+
+      def foo(self, x, y=1):
+        pass
+
+    constraint = function_type.FunctionType.from_callable(MyClass().foo)
+    self.assertEqual(
+        constraint,
+        function_type.FunctionType(
+            (function_type.Parameter(
+                "x", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                None),
+             function_type.Parameter(
+                 "y", function_type.Parameter.POSITIONAL_OR_KEYWORD, True,
+                 None))))
+    self.assertEqual(
+        function_type.FunctionType.get_default_values(MyClass().foo), {"y": 1})
+
+  def test_method_unbound(self):
+
+    class MyClass:
+
+      def foo(self, x, y=1):
+        pass
+
+    constraint = function_type.FunctionType.from_callable(MyClass.foo)
+    self.assertEqual(
+        constraint,
+        function_type.FunctionType(
+            (function_type.Parameter(
+                "self", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                None),
+             function_type.Parameter(
+                 "x", function_type.Parameter.POSITIONAL_OR_KEYWORD, False,
+                 None),
+             function_type.Parameter(
+                 "y", function_type.Parameter.POSITIONAL_OR_KEYWORD, True,
+                 None))))
+    self.assertEqual(
+        function_type.FunctionType.get_default_values(MyClass.foo), {"y": 1})
+
+  def test_required_only_validation(self):
+
+    def foo(x, y):  # pylint: disable=unused-argument
+      pass
+
+    constraint = function_type.FunctionType.from_callable(foo)
+    constraint.bind(*(1, 2))
+    constraint.bind(*(), **{"x": 1, "y": 2})
+    constraint.bind(*(), **{"y": 1, "x": 2})
+    constraint.bind(*(1,), **{"y": 2})
+
+    with self.assertRaisesRegex(TypeError, "too many positional arguments"):
+      constraint.bind(*(1, 2, 3))
+
+    with self.assertRaisesRegex(TypeError, "multiple values for argument 'x'"):
+      constraint.bind(*(1,), **{"x": 2})
+
+    with self.assertRaisesRegex(TypeError,
+                                "got an unexpected keyword argument 'z'"):
+      constraint.bind(*(1, 2), **{"z": 2})
+
+    with self.assertRaisesRegex(TypeError, "missing a required argument: 'x'"):
+      constraint.bind(*(), **{"z": 3})
+
+    with self.assertRaisesRegex(TypeError, "missing a required argument: 'x'"):
+      constraint.bind(*(), **{"y": 3})
+
+  def test_optional_only_validation(self):
+
+    def foo(x=1, y=2):  # pylint: disable=unused-argument
+      pass
+
+    constraint = function_type.FunctionType.from_callable(foo)
+    constraint.bind(*(1, 2))
+    constraint.bind(*(), **{"x": 1, "y": 2})
+    constraint.bind(*(1,), **{"y": 2})
+    constraint.bind(*(1,))
+    constraint.bind(*(), **{"x": 1})
+
+    with self.assertRaisesRegex(TypeError, "too many positional arguments"):
+      constraint.bind(*(1, 2, 3))
+
+    with self.assertRaisesRegex(TypeError, "multiple values for argument 'x'"):
+      constraint.bind(*(1,), **{"x": 2})
+
+    with self.assertRaisesRegex(TypeError,
+                                "got an unexpected keyword argument 'z'"):
+      constraint.bind(*(1, 2), **{"z": 2})
+
+    with self.assertRaisesRegex(TypeError,
+                                "got an unexpected keyword argument 'z'"):
+      constraint.bind(*(), **{"z": 3})
+
+  def test_required_and_optional_validation(self):
+
+    def foo(x, y=2):  # pylint: disable=unused-argument
+      pass
+
+    constraint = function_type.FunctionType.from_callable(foo)
+    constraint.bind(*(1, 2))
+    constraint.bind(*(), **{"x": 1, "y": 2})
+    constraint.bind(*(1,), **{"y": 2})
+    constraint.bind(*(1,))
+    constraint.bind(*(), **{"x": 1})
+
+    with self.assertRaisesRegex(TypeError, "too many positional arguments"):
+      constraint.bind(*(1, 2, 3))
+
+    with self.assertRaisesRegex(TypeError, "multiple values for argument 'x'"):
+      constraint.bind(*(1,), **{"x": 2})
+
+    with self.assertRaisesRegex(TypeError,
+                                "got an unexpected keyword argument 'z'"):
+      constraint.bind(*(1, 2), **{"z": 2})
+
+    with self.assertRaisesRegex(TypeError, "missing a required argument: 'x'"):
+      constraint.bind(*(), **{"z": 3})
+
+  def test_pickle(self):
+    original = function_type.FunctionType([
+        function_type.Parameter("x",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, None),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, None),
+        function_type.Parameter("z", function_type.Parameter.KEYWORD_ONLY,
+                                False, None)
+    ])
+    cloned = pickle.loads(pickle.dumps(original))
+    self.assertEqual(original, cloned)
+
+
+class CanonicalizationTest(test.TestCase, parameterized.TestCase):
+  args_1_2_3 = {"args": (1, 2, 3), "kwargs": {}}
+  args_1_2_kwargs_z_3 = {"args": (1, 2), "kwargs": {"z": 3}}
+  kwargs_x_1_y_2_z_3 = {"args": (), "kwargs": {"x": 1, "y": 2, "z": 3}}
+  args_1_2 = {"args": (1, 2), "kwargs": {}}
+  args_1_kwargs_y_2 = {"args": (1,), "kwargs": {"y": 2}}
+  kwargs_x_1_y_2 = {"args": (), "kwargs": {"x": 1, "y": 2}}
+
+  @parameterized.parameters(
+      args_1_2_3,
+      args_1_2_kwargs_z_3,
+      kwargs_x_1_y_2_z_3,
+  )
+  def test_required_only(self, args, kwargs):
+
+    def foo(x, y, z):
+      del x, y, z
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type)
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2, 3))
+    self.assertEqual(bound_args.kwargs, {})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  @parameterized.parameters(
+      args_1_2_3,
+      args_1_2_kwargs_z_3,
+      kwargs_x_1_y_2_z_3,
+  )
+  def test_optional_all(self, args, kwargs):
+
+    def foo(x=1, y=2, z=3):
+      del x, y, z
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type)
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2, 3))
+    self.assertEqual(bound_args.kwargs, {})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  @parameterized.parameters(args_1_2, args_1_kwargs_y_2, kwargs_x_1_y_2)
+  def test_optional_some(self, args, kwargs):
+
+    def foo(x=1, y=2, z=3):
+      del x, y, z
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    default_values = function_type.FunctionType.get_default_values(foo)
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, default_values, {},
+        polymorphic_type)
+
+    kwargs["z"] = 3
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2, 3))
+    self.assertEqual(bound_args.kwargs, {})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  @parameterized.parameters(
+      args_1_2_3,
+      args_1_2_kwargs_z_3,
+      kwargs_x_1_y_2_z_3,
+  )
+  def test_mixed(self, args, kwargs):
+
+    def foo(x, y, z=3):
+      del x, y, z
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type)
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2, 3))
+    self.assertEqual(bound_args.kwargs, {})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  def test_varargs(self):
+
+    def foo(*my_var_args):
+      del my_var_args
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    args = (1, 2, 3)
+    kwargs = {}
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type)
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2, 3))
+    self.assertEqual(bound_args.kwargs, {})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("my_var_args_0",
+                                function_type.Parameter.POSITIONAL_ONLY, False,
+                                trace_type.from_value(1, type_context)),
+        function_type.Parameter("my_var_args_1",
+                                function_type.Parameter.POSITIONAL_ONLY, False,
+                                trace_type.from_value(2, type_context)),
+        function_type.Parameter("my_var_args_2",
+                                function_type.Parameter.POSITIONAL_ONLY, False,
+                                trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  def test_varkwargs(self):
+
+    def foo(**kwargs):
+      del kwargs
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+
+    args = ()
+    kwargs = {"x": 1, "y": 2, "z": 3}
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type
+    )
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, ())
+    self.assertEqual(bound_args.kwargs, {"x": 1, "y": 2, "z": 3})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  def test_varargs_and_varkwargs(self):
+
+    def foo(*args, **kwargs):
+      del args, kwargs
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    args = (1,)
+    kwargs = {"y": 2, "z": 3}
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type
+    )
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1,))
+    self.assertEqual(bound_args.kwargs, {"y": 2, "z": 3})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("args_0",
+                                function_type.Parameter.POSITIONAL_ONLY, False,
+                                trace_type.from_value(1, type_context)),
+        function_type.Parameter("y", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  @parameterized.parameters(
+      args_1_2_kwargs_z_3,
+      kwargs_x_1_y_2_z_3,
+  )
+  def test_kwonly(self, args, kwargs):
+
+    def foo(x, y, *, z):
+      del x, y, z
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type)
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2))
+    self.assertEqual(bound_args.kwargs, {"z": 3})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+  @parameterized.parameters(
+      args_1_2_3,
+      args_1_2_kwargs_z_3,
+  )
+  def test_posonly(self, args, kwargs):
+    major, minor, _ = platform.python_version_tuple()
+    if not (major == "3" and int(minor) >= 8):
+      self.skipTest("Positional only args are supported in Python 3.8+")
+
+    # Raises syntax error in 3.7 but is important coverage for 3.8+.
+    foo = eval("lambda x, y, /, z: x + y + z")  # pylint: disable=eval-used
+
+    polymorphic_type = function_type.FunctionType.from_callable(foo)
+    mono_type, _ = function_type.canonicalize_to_monomorphic(
+        args, kwargs, {}, {}, polymorphic_type)
+
+    bound_args = mono_type.bind(*args, **kwargs)
+    self.assertEqual(bound_args.args, (1, 2, 3))
+    self.assertEqual(bound_args.kwargs, {})
+
+    type_context = trace_type.InternalTracingContext()
+    expected_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+
+    self.assertEqual(mono_type, expected_type)
+
+
+class TypeHierarchyTest(test.TestCase):
+
+  def test_same_type(self):
+    foo_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(1))
+    ])
+    self.assertEqual(foo_type, foo_type)
+    self.assertTrue(foo_type.is_supertype_of(foo_type))
+    self.assertEqual(
+        foo_type,
+        foo_type.most_specific_common_subtype([foo_type, foo_type, foo_type]))
+
+  def test_unrelated_types(self):
+    foo_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(1))
+    ])
+    bar_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(2))
+    ])
+    self.assertNotEqual(foo_type, bar_type)
+    self.assertFalse(foo_type.is_supertype_of(bar_type))
+    self.assertIsNone(
+        foo_type.most_specific_common_subtype([bar_type, bar_type]))
+    self.assertIsNone(
+        foo_type.most_specific_common_subtype([bar_type, foo_type]))
+
+  def test_partial_raises_error(self):
+    foo_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(1)),
+    ])
+    bar_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, None)
+    ])
+    self.assertNotEqual(foo_type, bar_type)
+
+    with self.assertRaises(TypeError):
+      foo_type.is_supertype_of(bar_type)
+
+    with self.assertRaises(TypeError):
+      bar_type.is_supertype_of(foo_type)
+
+    with self.assertRaises(TypeError):
+      foo_type.most_specific_common_subtype([bar_type, bar_type])
+
+    with self.assertRaises(TypeError):
+      bar_type.most_specific_common_subtype([foo_type, bar_type])
+
+  def test_related_types(self):
+
+    class MockAlwaysSuperType(trace.TraceType):
+
+      def is_subtype_of(self, other: trace.TraceType) -> bool:
+        return False
+
+      def most_specific_common_supertype(self, others):
+        return self
+
+      def placeholder_value(self, placeholder_context):
+        raise NotImplementedError
+
+      def __eq__(self, other):
+        return self is other
+
+      def __hash__(self):
+        return 0
+
+    supertype = MockAlwaysSuperType()
+
+    class MockAlwaysSubtype(trace.TraceType):
+
+      def is_subtype_of(self, other) -> bool:
+        return True
+
+      def most_specific_common_supertype(self, others):
+        return supertype
+
+      def placeholder_value(self, placeholder_context):
+        raise NotImplementedError
+
+      def __eq__(self, other):
+        return self is other
+
+      def __hash__(self):
+        return 1
+
+    subtype = MockAlwaysSubtype()
+
+    foo_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, supertype),
+    ])
+    bar_type = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, subtype)
+    ])
+
+    self.assertNotEqual(foo_type, bar_type)
+    self.assertTrue(bar_type.is_supertype_of(foo_type))
+    self.assertFalse(foo_type.is_supertype_of(bar_type))
+
+    self.assertEqual(
+        foo_type.most_specific_common_subtype([bar_type, foo_type]), foo_type)
+    self.assertEqual(
+        bar_type.most_specific_common_subtype([bar_type, foo_type]), foo_type)
+
+  def test_placeholder_arg(self):
+    type_context = trace_type.InternalTracingContext()
+    foo = function_type.FunctionType([
+        function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                                False, trace_type.from_value(1, type_context)),
+        function_type.Parameter("y",
+                                function_type.Parameter.POSITIONAL_OR_KEYWORD,
+                                False, trace_type.from_value(2, type_context)),
+        function_type.Parameter("z", function_type.Parameter.KEYWORD_ONLY,
+                                False, trace_type.from_value(3, type_context)),
+    ])
+    context_graph = func_graph.FuncGraph("test")
+    placeholder_context = trace_type.InternalPlaceholderContext(context_graph)
+    self.assertEqual(
+        foo.placeholder_arguments(placeholder_context).args, (1, 2))
+    self.assertEqual(
+        foo.placeholder_arguments(placeholder_context).kwargs, {"z": 3})
+
+
+class CapturesTest(test.TestCase):
+
+  def setUp(self):
+    super(CapturesTest, self).setUp()
+
+    def gen_type_fn(mapping):
+      return function_type.FunctionType([], collections.OrderedDict(mapping))
+
+    self.type_a1_b1 = gen_type_fn({
+        "a": trace_type.from_value(1),
+        "b": trace_type.from_value(1)
+    })
+    self.type_a1_b1_c1 = gen_type_fn({
+        "a": trace_type.from_value(1),
+        "b": trace_type.from_value(1),
+        "c": trace_type.from_value(1)
+    })
+    self.type_a2_b2_c2 = gen_type_fn({
+        "a": trace_type.from_value(2),
+        "b": trace_type.from_value(2),
+        "c": trace_type.from_value(2)
+    })
+    self.type_a1_b1_c2 = gen_type_fn({
+        "a": trace_type.from_value(1),
+        "b": trace_type.from_value(1),
+        "c": trace_type.from_value(2)
+    })
+    self.type_d1 = gen_type_fn({"d": trace_type.from_value(1)})
+
+  def testCapturesSubtype(self):
+    self.assertTrue(self.type_a1_b1.is_supertype_of(self.type_a1_b1_c1))
+    self.assertFalse(self.type_a1_b1_c1.is_supertype_of(self.type_a1_b1))
+    self.assertFalse(self.type_a1_b1_c1.is_supertype_of(self.type_a2_b2_c2))
+    self.assertFalse(self.type_a1_b1_c1.is_supertype_of(self.type_a2_b2_c2))
+    self.assertFalse(self.type_d1.is_supertype_of(self.type_a1_b1))
+
+  def testCapturesSupertype(self):
+    supertype_1 = self.type_a1_b1_c1.most_specific_common_subtype(
+        [self.type_a1_b1_c1])
+    self.assertLen(supertype_1.captures, 3)
+
+    supertype_2 = self.type_a1_b1.most_specific_common_subtype(
+        [self.type_a1_b1_c1, self.type_a2_b2_c2])
+    self.assertIsNone(supertype_2)
+
+    supertype_3 = self.type_a1_b1.most_specific_common_subtype(
+        [self.type_a1_b1_c2])
+    self.assertLen(supertype_3.captures, 3)
+
+    supertype_4 = self.type_a1_b1_c1.most_specific_common_subtype(
+        [self.type_a1_b1_c2])
+    self.assertIsNone(supertype_4)
+
+    supertype_5 = self.type_a1_b1_c1.most_specific_common_subtype(
+        [self.type_d1])
+    self.assertLen(supertype_5.captures, 4)
+
+
+class SanitizationTest(test.TestCase):
+
+  def testRename(self):
+    self.assertEqual("arg_42", function_type.sanitize_arg_name("42"))
+    self.assertEqual("a42", function_type.sanitize_arg_name("a42"))
+    self.assertEqual("arg__42", function_type.sanitize_arg_name("_42"))
+    self.assertEqual("a___", function_type.sanitize_arg_name("a%$#"))
+    self.assertEqual("arg____", function_type.sanitize_arg_name("%$#"))
+    self.assertEqual("foo", function_type.sanitize_arg_name("foo"))
+    self.assertEqual("Foo", function_type.sanitize_arg_name("Foo"))
+    self.assertEqual("arg_96ab_cd___53",
+                     function_type.sanitize_arg_name("96ab.cd//?53"))
+
+  def testLogWarning(self):
+
+    with self.assertLogs(level="WARNING") as logs:
+      result = function_type.sanitize_arg_name("96ab.cd//?53")
+
+    self.assertEqual(result, "arg_96ab_cd___53")
+
+    expected_message = (
+        "WARNING:absl:`96ab.cd//?53` is not a valid tf.function parameter name."
+        " Sanitizing to `arg_96ab_cd___53`.")
+    self.assertIn(expected_message, logs.output)
+
+
+class SerializationTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.product(
+      name=["arg_0", "param"],
+      kind=[
+          function_type.Parameter.POSITIONAL_ONLY,
+          function_type.Parameter.POSITIONAL_OR_KEYWORD
+      ],
+      optional=[True, False],
+      type_contraint=[None, trace_type.from_value(1)])
+  def testParameter(self, name, kind, optional, type_contraint):
+    original = function_type.Parameter(name, kind, optional, type_contraint)
+    expected_type_constraint = serialization.serialize(
+        type_contraint) if type_contraint else None
+    expected = function_type_pb2.Parameter(
+        name=name,
+        kind=function_type.PY_TO_PROTO_ENUM[kind],
+        is_optional=optional,
+        type_constraint=expected_type_constraint)
+    self.assertEqual(original.to_proto(), expected)
+    self.assertEqual(function_type.Parameter.from_proto(expected), original)
+
+  def testFunctionType(self):
+    original = function_type.FunctionType([
+        function_type.Parameter("a", function_type.Parameter.POSITIONAL_ONLY,
+                                False, None),
+    ], collections.OrderedDict([("b", trace_type.from_value(1))]))
+    expected = function_type_pb2.FunctionType(
+        parameters=[
+            function_type_pb2.Parameter(
+                name="a",
+                kind=function_type_pb2.Parameter.Kind.POSITIONAL_ONLY,
+                is_optional=False)
+        ],
+        captures=[
+            function_type_pb2.Capture(
+                name="b",
+                type_constraint=serialization.serialize(
+                    trace_type.from_value(1)))
+        ])
+    self.assertEqual(original.to_proto(), expected)
+    self.assertEqual(function_type.FunctionType.from_proto(expected), original)
+
+  def testCapturedDefaultValueStr(self):
+    f_type = function_type.FunctionType([
+        function_type.Parameter(
+            "a", function_type.Parameter.POSITIONAL_OR_KEYWORD, True, None
+        ),
+    ])
+
+    self.assertEqual(str(f_type), "(a=<captured_default_value>)")
+
+
+class FromStructuredSignatureTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      {
+          "signature": ((1, 2, 3), {}),
+          "expected_types": (
+              trace_type.from_value(1),
+              trace_type.from_value(2),
+              trace_type.from_value(3),
+          ),
+      },
+      {
+          "signature": (([1, 2, 3],), {}),
+          "expected_types": (
+              trace_type.from_value([1, 2, 3]),
+          ),
+      },
+      {
+          "signature": ((), {}),
+          "expected_types": (),
+      },
+  )
+  def testArgs(self, signature, expected_types):
+    generated_type = function_type.from_structured_signature(signature)
+    self.assertEqual(generated_type.output, trace_type.from_value(None))
+    for i, p in enumerate(generated_type.parameters.values()):
+      self.assertEqual(p.kind, function_type.Parameter.POSITIONAL_ONLY)
+      self.assertEqual(p.type_constraint, expected_types[i])
+
+  @parameterized.parameters(
+      {
+          "signature": ((), {"a": 1, "b": 2, "c": 3}),
+          "expected_types": {
+              "a": trace_type.from_value(1),
+              "b": trace_type.from_value(2),
+              "c": trace_type.from_value(3),
+          },
+      },
+      {
+          "signature": ((), {"a": [1, 2, 3]}),
+          "expected_types": {
+              "a": trace_type.from_value([1, 2, 3]),
+          },
+      },
+      {
+          "signature": ((), {}),
+          "expected_types": {},
+      },
+  )
+  def testKwargs(self, signature, expected_types):
+    generated_type = function_type.from_structured_signature(signature)
+    self.assertEqual(generated_type.output, trace_type.from_value(None))
+    for p in generated_type.parameters.values():
+      self.assertEqual(p.kind, function_type.Parameter.KEYWORD_ONLY)
+      self.assertEqual(p.type_constraint, expected_types[p.name])
+
+  @parameterized.parameters(
+      {"output_signature": 1},
+      {"output_signature": [1, 2, 3]},
+      {"output_signature": ()},
+  )
+  def testOutput(self, output_signature):
+    generated_type = function_type.from_structured_signature(
+        ((), {}), output_signature
+    )
+    self.assertEqual(
+        generated_type.output,
+        trace_type.from_value(
+            output_signature,
+            trace_type.InternalTracingContext(is_legacy_signature=True),
+        )
+    )
+
+if __name__ == "__main__":
+  test.main()

tensorflow/core/function/polymorphism/type_dispatch.py

@@ -0,0 +1,131 @@
+# 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.
+# ==============================================================================
+"""Polymorphic Type Dispatch."""
+
+import collections
+from typing import Optional, Iterable
+
+from tensorflow.core.function.polymorphism import function_type
+
+# The maximum number of dispatch lookups to cache.
+_MAX_DISPATCH_CACHE = 1024
+
+
+class TypeDispatchTable:
+  """Type dispatch table implementation.
+
+  A type dispatch table is a list, L, of target types. Given a request type, R,
+  the table selects a target type, T, according to the following dispatch rules:
+    1. R == T or R is supertype of T (functions are contravariant on args)
+    2. There does not exist O in L such that R is supertype of O and O is a
+       supertype of T (in other words, T is the closest to R, within list L).
+    3. If the above two rules are satisfied by multiple targets, the earliest
+       inserted one is chosen.
+  """
+
+  def __init__(self):
+    """Creates a TypeDispatchTable object."""
+    # Holds all inserted types as keys mapping to None.
+    # (Using OrderedDict as a set for determinism)
+    self._dispatch_table = collections.OrderedDict()
+
+    # LRU cache for dispatch results.
+    # Maps request types to target types (see class description).
+    # Does not contain exact matches, i.e, if cache[a] is b then a is not b.
+    self._dispatch_cache = collections.OrderedDict()
+
+  def add_target(self, target: function_type.FunctionType) -> None:
+    """Adds a new target type."""
+    self._dispatch_table[target] = None
+    for request in self._dispatch_cache:
+      if target.is_supertype_of(self._dispatch_cache[request]):
+        self._dispatch_cache[request] = target
+
+  @property
+  def targets(self) -> Iterable[function_type.FunctionType]:
+    """Returns an iterable to all targets in the table."""
+    return self._dispatch_table.keys()
+
+  def delete(self, target: function_type.FunctionType) -> None:
+    """Deletes a target in the table if it exists."""
+    if target in self._dispatch_table:
+      del self._dispatch_table[target]
+      for request in list(self._dispatch_cache.keys()):
+        if self._dispatch_cache[request] == target:
+          del self._dispatch_cache[request]
+
+  # TODO(b/205971333): remove once FunctionCache 'clear' is removed.
+  def clear(self) -> None:
+    """Deletes all targets in the table."""
+    self._dispatch_table.clear()
+    self._dispatch_cache.clear()
+
+  def dispatch(
+      self, request: function_type.FunctionType
+  ) -> Optional[function_type.FunctionType]:
+    """Returns the most specific supertype target if it exists in the table."""
+    # For known exact matches.
+    if request in self._dispatch_table:
+      return request
+
+    # For known non-exact matches.
+    # (self._dispatch cache does not contain exact matches)
+    if request in self._dispatch_cache:
+      # Move to the front of LRU cache.
+      result = self._dispatch_cache.pop(request)
+      self._dispatch_cache[request] = result
+      return result
+
+    most_specific_supertype = None
+    for other in self._dispatch_table:
+      if request.is_supertype_of(other):
+        if most_specific_supertype is None or other.is_supertype_of(
+            most_specific_supertype):
+          most_specific_supertype = other
+
+    self._cache_dispatch(request, most_specific_supertype)
+    return most_specific_supertype
+
+  def _cache_dispatch(self, request, target):
+    """Caches the dispatch lookup result for a target."""
+    if target is not None:
+      # LRU Cache removes oldest item
+      if len(self._dispatch_cache) > _MAX_DISPATCH_CACHE:
+        self._dispatch_cache.popitem(last=False)
+      self._dispatch_cache[request] = target
+
+  def try_generalizing_function_type(
+      self, target: function_type.FunctionType) -> function_type.FunctionType:
+    """Returns a generalized subtype of the one given.
+
+    This heuristic aims to reduce the number of future traces by computing a
+    type that represents more general function inputs.
+
+    The original "experimental_relax_shapes" heuristic identified a known type
+    which shared a common subtype with the current unknown type and then
+    traced with that common subtype. However, the notion of "common subtype"
+    was only limited to shapes. This heuristic extends that to FunctionType.
+
+    Returns `target` if a generalized subtype can not be found.
+
+    Args:
+      target: The FunctionType to generalize
+    """
+    relaxed = target
+    for other in self._dispatch_table:
+      subtype = relaxed.most_specific_common_subtype([other])
+      if subtype is not None:
+        relaxed = subtype
+    return relaxed

tensorflow/core/function/polymorphism/type_dispatch_test.py

@@ -0,0 +1,317 @@
+# 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 type_dispatch."""
+
+from typing import Optional
+
+from tensorflow.core.function.polymorphism import function_type
+from tensorflow.core.function.polymorphism import type_dispatch
+from tensorflow.python.platform import test
+from tensorflow.python.types import trace
+
+
+class MockShape(trace.TraceType):
+
+  def __init__(self, *shape: Optional[int]):
+    self.shape = shape
+
+  def is_subtype_of(self, other: "MockShape") -> bool:
+    if len(self.shape) != len(other.shape):
+      return False
+
+    return all(o is None or s == o for s, o in zip(self.shape, other.shape))
+
+  def most_specific_common_supertype(self, others):
+    if any(len(other.shape) != len(self.shape) for other in others):
+      return None
+
+    dims = [
+        dim if all(dim == other.shape[i]
+                   for other in others) else None
+        for i, dim in enumerate(self.shape)
+    ]
+    return MockShape(*dims)
+
+  def placeholder_value(self, placeholder_context=None):
+    raise NotImplementedError
+
+  def __str__(self):
+    return str(self.shape)
+
+  def __repr__(self):
+    return str(self)
+
+  def __hash__(self) -> int:
+    return hash(self.shape)
+
+  def __eq__(self, other: "MockShape") -> bool:
+    return self.shape == other.shape
+
+
+def make_shape_function_type(*shape):
+  return function_type.FunctionType([
+      function_type.Parameter("x", function_type.Parameter.POSITIONAL_ONLY,
+                              False, MockShape(*shape))
+  ])
+
+
+class TypeDispatchTableTest(test.TestCase):
+
+  def testVertical(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None, None))
+    table.add_target(make_shape_function_type(None, None, 1))
+    table.add_target(make_shape_function_type(None, 1, 1))
+    table.add_target(make_shape_function_type(1, 1, 1))
+    self.assertEqual(
+        list(table.targets), [
+            make_shape_function_type(None, None, None),
+            make_shape_function_type(None, None, 1),
+            make_shape_function_type(None, 1, 1),
+            make_shape_function_type(1, 1, 1)
+        ])
+
+  def testHorizontal(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(1,))
+    table.add_target(make_shape_function_type(1, 2))
+    table.add_target(make_shape_function_type(1, 2, 3))
+    self.assertEqual(
+        list(table.targets), [
+            make_shape_function_type(1,),
+            make_shape_function_type(1, 2),
+            make_shape_function_type(1, 2, 3)
+        ])
+
+  def testDuplicateNodes(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None))
+    table.add_target(make_shape_function_type(1, None))
+    table.add_target(make_shape_function_type(None, 2))
+    table.add_target(make_shape_function_type(None, None))
+    self.assertEqual(
+        list(table.targets), [
+            make_shape_function_type(None, None),
+            make_shape_function_type(1, None),
+            make_shape_function_type(None, 2)
+        ])
+
+  def testDeletion(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None))
+    table.add_target(make_shape_function_type(None, 1))
+    table.add_target(make_shape_function_type(None, 2))
+
+    self.assertEqual(
+        list(table.targets), [
+            make_shape_function_type(None, None),
+            make_shape_function_type(None, 1),
+            make_shape_function_type(None, 2)
+        ])
+
+    table.delete(make_shape_function_type(None, 2))  # Should remove the target
+
+    self.assertEqual(
+        list(table.targets), [
+            make_shape_function_type(None, None),
+            make_shape_function_type(None, 1),
+        ])
+
+    table.delete(make_shape_function_type(None, 2))  # Should have no effect
+
+    self.assertEqual(
+        list(table.targets), [
+            make_shape_function_type(None, None),
+            make_shape_function_type(None, 1),
+        ])
+
+  def testContains(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None, None))
+    table.add_target(make_shape_function_type(None, 1))
+    table.add_target(make_shape_function_type(1, 1))
+    table.add_target(make_shape_function_type(None, 2, 1))
+
+    self.assertIn(make_shape_function_type(None, None, None), table.targets)
+    self.assertIn(make_shape_function_type(None, 1), table.targets)
+    self.assertIn(make_shape_function_type(1, 1), table.targets)
+    self.assertIn(make_shape_function_type(None, 2, 1), table.targets)
+
+    self.assertNotIn(make_shape_function_type(None, None, 1), table.targets)
+    self.assertNotIn(make_shape_function_type(1, None), table.targets)
+    self.assertNotIn(make_shape_function_type(1, 2), table.targets)
+    self.assertNotIn(make_shape_function_type(None, 2, None), table.targets)
+
+  def testDispatchExactMatches(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None, None))
+    table.add_target(make_shape_function_type(None, 1, None))
+    table.add_target(make_shape_function_type(None, 1, 2))
+    table.add_target(make_shape_function_type(None, 2, 2))
+
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(None, 1, 2)),
+        make_shape_function_type(None, 1, 2))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(None, 1, None)),
+        make_shape_function_type(None, 1, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(None, None, None)),
+        make_shape_function_type(None, None, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(None, 2, 2)),
+        make_shape_function_type(None, 2, 2))
+
+  def testDispatchMoreSpecific(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None, None))
+    table.add_target(make_shape_function_type(None, 1, None))
+    table.add_target(make_shape_function_type(None, 1, 2))
+    table.add_target(make_shape_function_type(None, 2, 2))
+
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, 1, 2))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 3)),
+        make_shape_function_type(None, 1, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 3, 3)),
+        make_shape_function_type(None, None, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 2, 2)),
+        make_shape_function_type(None, 2, 2))
+
+  def testDispatchNoMatches(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, 1, None))
+    table.add_target(make_shape_function_type(None, 1, 2))
+    table.add_target(make_shape_function_type(None, 2, 2))
+
+    self.assertIsNone(table.dispatch(make_shape_function_type(1, 2)))
+    self.assertIsNone(table.dispatch(make_shape_function_type(1, 2, 3)))
+    self.assertIsNone(table.dispatch(make_shape_function_type(1, 2, 3, 4)))
+
+  def testDispatchCachedAddUpdates(self):
+    table = type_dispatch.TypeDispatchTable()
+
+    table.add_target(make_shape_function_type(None, None, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, None, None))
+
+    table.add_target(make_shape_function_type(None, 1, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, 1, None))
+
+    table.add_target(make_shape_function_type(None, 1, 2))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, 1, 2))
+
+    table.add_target(make_shape_function_type(1, 1, 2))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(1, 1, 2))
+
+  def testDispatchCachedDeleteUpdates(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None, None))
+    table.add_target(make_shape_function_type(None, 1, None))
+    table.add_target(make_shape_function_type(None, 1, 2))
+    table.add_target(make_shape_function_type(1, 1, 2))
+
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(1, 1, 2))
+
+    table.delete(make_shape_function_type(1, 1, 2))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, 1, 2))
+
+    table.delete(make_shape_function_type(None, 1, 2))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, 1, None))
+
+    table.delete(make_shape_function_type(None, 1, None))
+    self.assertEqual(
+        table.dispatch(make_shape_function_type(1, 1, 2)),
+        make_shape_function_type(None, None, None))
+
+  def testDispatchCacheOrderingDeterminism(self):
+    table_1 = type_dispatch.TypeDispatchTable()
+    table_1.add_target(make_shape_function_type(1, None, None))
+    table_1.add_target(make_shape_function_type(None, 2, None))
+    table_1.add_target(make_shape_function_type(None, None, 3))
+
+    table_2 = type_dispatch.TypeDispatchTable()
+    table_2.add_target(make_shape_function_type(None, 2, None))
+    table_2.add_target(make_shape_function_type(1, None, None))
+    table_2.add_target(make_shape_function_type(None, None, 3))
+
+    table_3 = type_dispatch.TypeDispatchTable()
+    table_3.add_target(make_shape_function_type(None, None, 3))
+    table_3.add_target(make_shape_function_type(1, None, None))
+    table_3.add_target(make_shape_function_type(None, 2, None))
+
+    # table_1, table_2, table_3 have the same targets
+    self.assertEqual(set(table_1.targets), set(table_2.targets))
+    self.assertEqual(set(table_2.targets), set(table_3.targets))
+
+    # But they dispatch to the first target they find which does not have any
+    # more specific viable target.
+    shape = make_shape_function_type(1, 2, 3)
+    self.assertEqual(
+        table_1.dispatch(shape), make_shape_function_type(1, None, None))
+    self.assertEqual(
+        table_2.dispatch(shape), make_shape_function_type(None, 2, None))
+    self.assertEqual(
+        table_3.dispatch(shape), make_shape_function_type(None, None, 3))
+
+  def testGeneralizedExisting(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, None, None))
+    table.add_target(make_shape_function_type(None, 1, None))
+    table.add_target(make_shape_function_type(None, 1, 2))
+    self.assertEqual(
+        table.try_generalizing_function_type(
+            make_shape_function_type(None, 1, 3)),
+        make_shape_function_type(None, None, None))
+
+  def testGeneralizedNovel(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, 1, None))
+    table.add_target(make_shape_function_type(None, 1, 2))
+    self.assertEqual(
+        table.try_generalizing_function_type(
+            make_shape_function_type(None, 2, 3)),
+        make_shape_function_type(None, None, None))
+
+  def testGeneralizedUnknown(self):
+    table = type_dispatch.TypeDispatchTable()
+    table.add_target(make_shape_function_type(None, 1))
+    table.add_target(make_shape_function_type(None, 2))
+    table.add_target(make_shape_function_type(None, 3))
+    self.assertEqual(
+        table.try_generalizing_function_type(
+            make_shape_function_type(None, 4, 3)),
+        make_shape_function_type(None, 4, 3))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/core/function/runtime_client/BUILD

@@ -0,0 +1,182 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension")
+load("//tensorflow/core/platform:build_config_root.bzl", "if_static")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/core/function:__subpackages__"],
+)
+
+licenses(["notice"])
+
+cc_library(
+    name = "runtime_client_cc",
+    srcs = [
+        "runtime_client.cc",
+    ],
+    hdrs = [
+        "runtime_client.h",
+    ],
+    defines = select({
+        "//tensorflow/compiler/mlir/python:disable_mlir_config": ["DISABLE_MLIR"],
+        "//conditions:default": [],
+    }),
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:immediate_execution_context",
+        "//tensorflow/c/eager:immediate_execution_operation",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:import_model",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/common_runtime:function_def_utils",
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/common_runtime/eager:core",
+        "//tensorflow/core/framework:function_proto_cc",
+        "//tensorflow/core/framework:graph_proto_cc",
+        "//tensorflow/core/framework:op_def_proto_cc",
+        "//tensorflow/core/ir:Dialect",
+        "//tensorflow/core/ir/importexport:graphdef_export",
+        "//tensorflow/core/ir/importexport:graphdef_import",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/platform:stringpiece",
+        "//tensorflow/core/platform:types",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ] + select({
+        "//tensorflow/compiler/mlir/python:disable_mlir_config": [],
+        "//conditions:default": [
+            "//tensorflow/compiler/mlir/python:mlir",
+        ],
+    }),
+    # TODO(mdan): Get rid of alwayslink, it's nonstandard.
+    alwayslink = 1,
+)
+
+# TODO(mdan): Pull these transitive header deps in a more decent fashion.
+# TODO(mdan): Get rid of headers-only lib, it's nonstandard. Use cc_shared_library?
+cc_library(
+    name = "runtime_client_headers",
+    textual_hdrs = [
+        "runtime_client.h",
+        "//tensorflow/c/eager:pywrap_required_hdrs",
+        "//tensorflow/core/common_runtime/eager:pywrap_required_hdrs",
+        "//tensorflow/core/config:flags_headers",
+        "//tensorflow/core/distributed_runtime:pywrap_required_hdrs",
+        "//tensorflow/core/distributed_runtime/coordination:pywrap_required_hdrs",
+        "//tensorflow/core/distributed_runtime/eager:pywrap_required_hdrs",
+        "@xla//xla/tsl/distributed_runtime:pywrap_required_hdrs",
+        "@xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs",
+    ],
+)
+
+tf_cc_test(
+    name = "runtime_client_cc_test",
+    srcs = ["runtime_client_test.cc"],
+    deps = [
+        ":runtime_client_cc",
+        "//tensorflow/c:c_api_experimental",  # buildcleaner: keep (registers CPU ops?)
+        "//tensorflow/c:tensor_interface",
+        "//tensorflow/c/eager:immediate_execution_tensor_handle",
+        "//tensorflow/core:test",
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/framework:function_proto_cc",
+        "//tensorflow/core/framework:op_def_proto_cc",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/function/testing:test_pass_cc",
+        "//tensorflow/core/ir:Dialect",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:statusor",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:Parser",
+    ],
+)
+
+# TODO(b/221223841): Get rid of if_static, it's nonstandard.
+tf_python_pybind_extension(
+    name = "runtime_client_pybind",
+    srcs = ["runtime_client_pybind.cc"],
+    data = [
+        "runtime_client_pybind.pyi",
+    ],
+    enable_stub_generation = True,
+    deps = [
+        ":runtime_client_headers",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@com_google_absl//absl/status",
+        "@pybind11",
+    ] + if_static(
+        extra_deps = [
+            "//tensorflow/core/framework:function_proto_cc",
+            "//tensorflow/core/protobuf:eager_service_proto_cc",
+            "//tensorflow/core/protobuf:master_proto_cc",
+            "//tensorflow/core/protobuf:worker_proto_cc",
+            "@xla//xla/tsl/protobuf:coordination_service_proto_cc",
+        ],
+        otherwise = [
+            "//tensorflow/core/framework:function_proto_cc_headers_only",
+            "//tensorflow/core/protobuf:eager_service_proto_cc_headers_only",
+            "//tensorflow/core/protobuf:master_proto_cc_headers_only",
+            "//tensorflow/core/protobuf:worker_proto_cc_headers_only",
+            "@xla//xla/tsl/protobuf:coordination_service_proto_cc_headers_only",
+        ],
+    ),
+)
+
+# TODO(mdan): Drop function_proto_py_pb2 once pybind11_protobuf is available.
+pytype_strict_library(
+    name = "runtime_client_py",
+    srcs = [
+        "runtime_client.py",
+    ],
+    visibility = [
+        "//learning/brain/experimental/tfq:__subpackages__",
+        "//tensorflow/core/function/transform:__subpackages__",
+        "//tensorflow/python/eager:__subpackages__",
+    ],
+    deps = [
+        ":runtime_client_pybind",
+        "//tensorflow/core/framework:function_proto_py",
+        "//tensorflow/python:pywrap_tensorflow",  # buildcleaner: keep (required for TF pybind)
+    ],
+)
+
+py_test(
+    name = "runtime_client_py_test",
+    srcs = ["runtime_client_test.py"],
+    main = "runtime_client_test.py",
+    strict_deps = True,
+    tags = ["no_oss"],  # TODO(b/219089812)
+    deps = [
+        ":runtime_client_py",
+        #internal proto upb dep
+        "//tensorflow/core/framework:function_proto_py",
+        "//tensorflow/core/function/testing:test_pass_py",
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:execute",
+        "//tensorflow/python/eager:remote",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/core/function/runtime_client/runtime_client.cc

@@ -0,0 +1,256 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/function/runtime_client/runtime_client.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/immediate_execution_context.h"
+#include "tensorflow/c/eager/immediate_execution_operation.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "tensorflow/core/graph/graph.h"
+
+#if !defined(DISABLE_MLIR)
+#include "tensorflow/compiler/mlir/python/mlir.h"
+#endif
+
+#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/common_runtime/function_def_utils.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/device_factory.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/ir/importexport/graphdef_export.h"
+#include "tensorflow/core/ir/importexport/graphdef_import.h"
+#include "tensorflow/core/ir/ops.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+namespace core {
+namespace function {
+
+EagerContext& GlobalEagerContext() {
+  static EagerContext* global_ctx = []() {
+    SessionOptions opts;
+    std::vector<std::unique_ptr<Device>> devices;
+    absl::Status&& device_init_status = DeviceFactory::AddDevices(
+        opts, "/job:localhost/replica:0/task:0", &devices);
+    CHECK(device_init_status.ok());  // Crash OK
+
+    return new EagerContext(
+        opts, ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
+        /*async=*/false,
+        /*device_mgr=*/new DynamicDeviceMgr(std::move(devices)),
+        /*device_mgr_owned=*/true,
+        /*rendezvous=*/nullptr,
+        /*cluster_flr=*/nullptr,
+        /*collective_executor_mgr=*/nullptr,
+        /*run_eager_op_as_function=*/true);
+  }();
+  return *global_ctx;
+}
+
+EagerContext& GlobalPythonEagerContext() {
+  EagerContext* ctx = reinterpret_cast<EagerContext*>(GetCEagerContext());
+  DCHECK(ctx) << "The Python eager context must be initialized first.";
+  return *ctx;
+}
+
+absl::StatusOr<FunctionDef> Runtime::GetFunctionProto(absl::string_view name) {
+  EagerContext& ctx = this->eager_ctx_;
+
+  const FunctionDef* f = ctx.FindFunctionDef(std::string(name));
+  if (f == nullptr) {
+    return absl::Status(
+        absl::StatusCode::kInvalidArgument,
+        absl::StrCat("Could not find an attribute for key ", name));
+  }
+
+  return *f;
+}
+
+absl::Status Runtime::CreateFunction(const FunctionDef& fdef) {
+  const auto& fname = fdef.signature().name();
+  if (this->eager_ctx_.FindFunctionByName(fname)) {
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(this->eager_ctx_.RemoveFunction(fname),
+                                    "removing function ", fname);
+  }
+  return this->eager_ctx_.AddFunctionDef(fdef);
+}
+
+absl::Status Runtime::CreateFunction(OpaqueTfgGraphFuncOp* fop) {
+  mlir::tfg::GraphFuncOp fop_proper =
+      *reinterpret_cast<mlir::tfg::GraphFuncOp*>(fop);
+  return mlir::tfg::ConvertToFunctionDef(fop_proper,
+                                         *this->eager_ctx_.FuncLibDef());
+}
+
+absl::Status Runtime::CreateFunction(OpaqueTfFuncOp* fop) {
+  mlir::func::FuncOp fop_proper = *reinterpret_cast<mlir::func::FuncOp*>(fop);
+  const auto& fname = fop_proper.getName().str();
+  GraphExportConfig config;
+  FunctionDef fdef;
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      tf2xla::v2::ConvertMlirFunctionToFunctionLibraryDef(fop_proper, config,
+                                                          &fdef),
+      "creating function ", fname);
+  return CreateFunction(fdef);
+}
+
+absl::Status Runtime::TransformFunction(absl::string_view name,
+                                        absl::string_view pipeline_name,
+                                        Dialect dialect) {
+  // TODO(mdan): Use a longer-lived context.
+  mlir::MLIRContext ctx;
+  mlir::PassManager pm(&ctx);
+
+  std::string error;
+  llvm::raw_string_ostream error_stream(error);
+  // StringPiece doesn't seem to always be compatible with StringRef.
+  if (mlir::failed(mlir::parsePassPipeline(std::string(pipeline_name), pm,
+                                           error_stream))) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        absl::StrCat("locating pass pipeline ", pipeline_name,
+                                     ": ", error_stream.str()));
+  }
+
+  // For now, we roundtrip from proto. Once we have a permanent MLIR
+  // representation, we should be able to use it directly.
+  auto fn = GetFunctionProto(name);
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(fn.status(), "loading function ", name);
+
+  GraphDef graph;
+  *graph.mutable_library()->add_function() = *fn;
+  tensorflow::GraphDebugInfo debug_info;
+  // TODO(xjun): Hoist branches into helper functions.
+  if (dialect == Dialect::TFG) {
+    auto mlir_fn = mlir::tfg::ImportGraphDef(&ctx, debug_info, graph);
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(mlir_fn.status(), "importing function ",
+                                    name);
+
+    mlir::StatusScopedDiagnosticHandler diagnostics_handler(&ctx);
+    if (failed(pm.run(mlir_fn->get()))) {
+      return diagnostics_handler.Combine(absl::Status(
+          absl::StatusCode::kInvalidArgument,
+          absl::StrCat("running pass pipeline ", pipeline_name, ": ")));
+    }
+
+    for (auto fn : mlir_fn->get().getBody()->getOps<mlir::tfg::GraphFuncOp>()) {
+      TF_RETURN_WITH_CONTEXT_IF_ERROR(
+          CreateFunction(reinterpret_cast<OpaqueTfgGraphFuncOp*>(&fn)),
+          absl::StrCat("updating function ", fn.getName().str()));
+    }
+    return absl::OkStatus();
+  }
+
+  if (dialect == Dialect::TF) {
+    absl::Status status;
+    FunctionLibraryDefinition& flib_def = *this->eager_ctx_.FuncLibDef();
+    std::unique_ptr<FunctionBody> fbody;
+    status = FunctionDefToBodyHelper(*fn, AttrSlice(), &flib_def, &fbody);
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(status, "importing function ", name);
+
+    tensorflow::GraphImportConfig specs;
+    specs.graph_func_name = fbody->record->fdef().signature().name();
+    specs.enable_shape_inference = false;
+    specs.graph_as_function = true;
+    for (const Node* control_ret_node : fbody->control_ret_nodes)
+      specs.control_outputs.push_back(control_ret_node->name());
+    absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> mlir_fn =
+        tensorflow::tf2xla::v2::ConvertGraphToTfExecutor(*fbody->graph, {},
+                                                         flib_def, specs, &ctx);
+    TF_RETURN_WITH_CONTEXT_IF_ERROR(mlir_fn.status(), "importing function ",
+                                    name);
+
+    mlir::StatusScopedDiagnosticHandler diagnostics_handler(&ctx);
+    if (failed(pm.run(mlir_fn->get()))) {
+      return diagnostics_handler.Combine(absl::Status(
+          absl::StatusCode::kInvalidArgument,
+          absl::StrCat("running pass pipeline ", pipeline_name, ": ")));
+    }
+
+    for (auto fn : mlir_fn->get().getBody()->getOps<mlir::func::FuncOp>()) {
+      TF_RETURN_WITH_CONTEXT_IF_ERROR(
+          CreateFunction(reinterpret_cast<OpaqueTfFuncOp*>(&fn)),
+          absl::StrCat("updating function ", fn.getName().str()));
+    }
+    return absl::OkStatus();
+  }
+
+  return absl::Status(
+      absl::StatusCode::kInvalidArgument,
+      absl::StrCat("Unsupported dialect: ", dialect,
+                   ". Supported dialects are Dialect::TFG and Dialect::TF."));
+}
+
+absl::StatusOr<ReturnValues> Runtime::CallFunction(
+    absl::string_view name, absl::Span<AbstractTensorHandle* const> args) {
+  EagerContext& ctx = this->eager_ctx_;
+
+  ImmediateOpPtr op(ctx.CreateOperation());
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(op->Reset(name.data(), nullptr),
+                                  "initializing call op for ", name);
+
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(op->AddInputList(args),
+                                  "preparing call args for ", name);
+
+  const FunctionDef* fn_def = ctx.GetFunctionDef(std::string(name));
+  int num_retvals = fn_def->signature().output_arg_size();
+  int actual_retvals = num_retvals;
+  std::vector<ImmediateExecutionTensorHandle*> retvals(num_retvals);
+  TF_RETURN_WITH_CONTEXT_IF_ERROR(
+      op->Execute(absl::MakeSpan(
+                      reinterpret_cast<AbstractTensorHandle**>(retvals.data()),
+                      num_retvals),
+                  &actual_retvals),
+      "executing call op for ", name);
+  DCHECK(num_retvals == actual_retvals);
+
+  ReturnValues final_returns;
+  for (const auto& r : retvals) {
+    final_returns.emplace_back(ImmediateTensorHandlePtr(r));
+  }
+
+  return final_returns;
+}
+
+}  // namespace function
+}  // namespace core
+}  // namespace tensorflow

tensorflow/core/function/runtime_client/runtime_client.h

@@ -0,0 +1,100 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FUNCTION_RUNTIME_CLIENT_RUNTIME_CLIENT_H_
+#define TENSORFLOW_CORE_FUNCTION_RUNTIME_CLIENT_RUNTIME_CLIENT_H_
+
+#include <vector>
+
+#include "absl/types/span.h"
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/platform/stringpiece.h"
+
+namespace tensorflow {
+namespace core {
+namespace function {
+
+// TODO(mdan): Get rid of this once pybind can depend on MLIR headers.
+// This empty struct serves to hide a pointer to an actual MLIR TFG dialect
+// FuncOp object.
+struct OpaqueTfgGraphFuncOp;
+
+// TODO(xjun): Get rid of this once pybind can depend on MLIR headers.
+// This empty struct serves to hide a pointer to an actual MLIR TF dialect
+// FuncOp object.
+struct OpaqueTfFuncOp;
+
+// This is the current global context managed by the Python API. For historical
+// reasons, the Python runtime controls this context and all other clients must
+// use it. See tensorflow/python/eager/pywrap_tfe.h and
+// tensorflow/python/eager/context.py.
+//
+// This must always be called after the Python eager context was initialized.
+//
+// If the Python runtime isn't involved, or when writing code that exclusively
+// relies on functions defined in this namespace, users are encouraged to
+// maintain their own EagerContext or use GlobalEagerContext.
+EagerContext& GlobalPythonEagerContext();
+
+// This global context is available for testing and to be shared among various
+// APIs.
+EagerContext& GlobalEagerContext();
+
+using ReturnValues = std::vector<ImmediateTensorHandlePtr>;
+
+// A public API for manipulating and executing functions in a TensorFlow
+// runtime.
+class Runtime {
+ public:
+  explicit Runtime(EagerContext& eager_ctx) : eager_ctx_(eager_ctx) {}
+
+  enum class Dialect {
+    TFG,
+    TF,
+  };
+
+  absl::StatusOr<FunctionDef> GetFunctionProto(absl::string_view name);
+
+  // TODO(mdan): Enforce creation or rename to SetFunction.
+  absl::Status CreateFunction(const FunctionDef& fdef);
+  // TODO(mdan): Change to mlir::tfg::GraphFuncOp once pybind can depend on it.
+  absl::Status CreateFunction(OpaqueTfgGraphFuncOp* fop);
+  // TODO(xjun): Change to mlir::func::FuncOp once pybind can depend on it.
+  absl::Status CreateFunction(OpaqueTfFuncOp* fop);
+  // Applies a MLIR pipeline to an existing function.
+  // The pipeline may rename the function. If it does so, the old function
+  // remains unchanged. If the new name specifies an existing function, it will
+  // be overwritten.
+  absl::Status TransformFunction(absl::string_view name,
+                                 absl::string_view pipeline_name,
+                                 Dialect dialect = Dialect::TFG);
+
+  absl::StatusOr<ReturnValues> CallFunction(
+      absl::string_view name, absl::Span<AbstractTensorHandle* const> args);
+
+ private:
+  EagerContext& eager_ctx_;
+};
+
+}  // namespace function
+}  // namespace core
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_FUNCTION_RUNTIME_CLIENT_RUNTIME_CLIENT_H_

tensorflow/core/function/runtime_client/runtime_client.py

@@ -0,0 +1,35 @@
+# 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.
+# ==============================================================================
+"""Low level TF runtime client."""
+
+# TF oddity: this import loads TF-specific dynamic libraries.
+from tensorflow.python import pywrap_tensorflow  # pylint:disable=g-bad-import-order,unused-import
+
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.function.runtime_client import runtime_client_pybind
+
+GlobalEagerContext = runtime_client_pybind.GlobalEagerContext
+GlobalPythonEagerContext = runtime_client_pybind.GlobalPythonEagerContext
+
+
+# TODO(mdan): Map without adapters once pybind11_protobuf available
+class Runtime(runtime_client_pybind.Runtime):
+
+  def GetFunctionProto(self, name: str) -> function_pb2.FunctionDef:
+    return function_pb2.FunctionDef.FromString(
+        self.GetFunctionProtoString(name))
+
+  def CreateFunction(self, function_def: function_pb2.FunctionDef):
+    self.CreateFunctionFromString(function_def.SerializeToString())

tensorflow/core/function/runtime_client/runtime_client_pybind.cc

@@ -0,0 +1,70 @@
+/* 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.
+==============================================================================*/
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "pybind11/attr.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/c/eager/abstract_tensor_handle.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/function/runtime_client/runtime_client.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+PYBIND11_MAKE_OPAQUE(tensorflow::EagerContext);
+
+PYBIND11_MODULE(runtime_client_pybind, m) {
+  pybind11::class_<tensorflow::EagerContext, tensorflow::EagerContextPtr>
+      EagerContext(m, "EagerContext");
+  pybind11::class_<absl::Status> Status(m, "Status", pybind11::module_local());
+
+  m.def("GlobalEagerContext", &tensorflow::core::function::GlobalEagerContext,
+        pybind11::return_value_policy::reference);
+
+  m.def("GlobalPythonEagerContext",
+        &tensorflow::core::function::GlobalPythonEagerContext,
+        pybind11::return_value_policy::reference);
+
+  pybind11::class_<tensorflow::core::function::Runtime> runtime(m, "Runtime");
+
+  pybind11::enum_<tensorflow::core::function::Runtime::Dialect>(runtime,
+                                                                "Dialect")
+      .value("TFG", tensorflow::core::function::Runtime::Dialect::TFG)
+      .value("TF", tensorflow::core::function::Runtime::Dialect::TF);
+
+  runtime.def(pybind11::init<tensorflow::EagerContext&>());
+  // TODO(mdan): Rename to GetFunctionProto once pybind11_protobuf available
+  runtime.def(
+      "GetFunctionProtoString",
+      [](tensorflow::core::function::Runtime& r, const std::string& name) {
+        return pybind11::bytes(r.GetFunctionProto(name)->SerializeAsString());
+      },
+      pybind11::return_value_policy::reference);
+  // TODO(mdan): Rename to CreateFunction once pybind11_protobuf available
+  runtime.def(
+      "CreateFunctionFromString",
+      [](tensorflow::core::function::Runtime& r, const std::string& def) {
+        tensorflow::FunctionDef proto;
+        proto.ParseFromString(def);
+        return r.CreateFunction(proto);
+      });
+  runtime.def("TransformFunction",
+              &tensorflow::core::function::Runtime::TransformFunction,
+              pybind11::arg("name"), pybind11::arg("pipeline_name"),
+              pybind11::arg("dialect") =
+                  tensorflow::core::function::Runtime::Dialect::TFG);
+}

tensorflow/core/function/runtime_client/runtime_client_pybind.pyi

@@ -0,0 +1,46 @@
+# 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 typing import ClassVar
+
+class EagerContext:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class Runtime:
+    class Dialect:
+        __members__: ClassVar[dict] = ...  # read-only
+        TF: ClassVar[Runtime.Dialect] = ...
+        TFG: ClassVar[Runtime.Dialect] = ...
+        __entries: ClassVar[dict] = ...
+        def __init__(self, value: int) -> None: ...
+        def __eq__(self, other: object) -> bool: ...
+        def __hash__(self) -> int: ...
+        def __index__(self) -> int: ...
+        def __int__(self) -> int: ...
+        def __ne__(self, other: object) -> bool: ...
+        @property
+        def name(self) -> str: ...
+        @property
+        def value(self) -> int: ...
+    def __init__(self, arg0: EagerContext) -> None: ...
+    def CreateFunctionFromString(self, arg0: str) -> Status: ...
+    def GetFunctionProtoString(self, arg0: str) -> bytes: ...
+    def TransformFunction(self, name: str, pipeline_name: str, dialect: Runtime.Dialect = ...) -> Status: ...
+
+class Status:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def GlobalEagerContext() -> EagerContext: ...
+def GlobalPythonEagerContext() -> EagerContext: ...

tensorflow/core/function/runtime_client/runtime_client_test.cc

@@ -0,0 +1,387 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/function/runtime_client/runtime_client.h"
+
+#include <stdint.h>
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/core/common_runtime/eager/context.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/function/testing/test_pass.h"
+#include "tensorflow/core/ir/dialect.h"
+#include "tensorflow/core/ir/ops.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace core {
+namespace function {
+namespace {
+
+EagerContextPtr TestingEagerCtx() {
+  SessionOptions opts;
+  std::vector<std::unique_ptr<Device>> devices;
+  absl::Status&& device_init_status = DeviceFactory::AddDevices(
+      opts, "/job:localhost/replica:0/task:0", &devices);
+  CHECK(device_init_status.ok());  // Crash OK
+
+  return EagerContextPtr(new EagerContext(
+      opts, ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
+      /*async=*/false,
+      /*device_mgr=*/new DynamicDeviceMgr(std::move(devices)),
+      /*device_mgr_owned=*/true,
+      /*rendezvous=*/nullptr,
+      /*cluster_flr=*/nullptr,
+      /*collective_executor_mgr=*/nullptr,
+      /*run_eager_op_as_function=*/true));
+}
+
+int IntValue(ImmediateExecutionTensorHandle& h) {
+  absl::Status status;
+  AbstractTensorPtr t(h.Resolve(&status));
+  DCHECK(status.ok());
+  switch (h.DataType()) {
+    case DT_INT32:
+      return *(static_cast<int32_t*>(t->Data()));
+    case DT_INT64:
+      return *(static_cast<int64_t*>(t->Data()));
+    default:
+      DCHECK(false) << "invalid data type";
+      return 0;
+  }
+}
+
+ImmediateTensorHandlePtr IntScalarTensor(EagerContext& ctx, int value) {
+  AbstractTensorPtr tensor(ctx.CreateInt32Scalar(value));
+  ImmediateTensorHandlePtr handle(ctx.CreateLocalHandle(tensor.get()));
+  return handle;
+}
+
+FunctionDef MakeNullaryFunction() {
+  FunctionDef fd;
+  protobuf::TextFormat::Parser parser;
+  CHECK(parser.ParseFromString(
+      R"pb(signature {
+             name: 'NullaryFunction'
+             output_arg { name: 'o' type: DT_INT32 }
+           }
+           node_def {
+             name: 'retval'
+             op: 'Const'
+             attr {
+               key: 'dtype'
+               value { type: DT_INT32 }
+             }
+             attr {
+               key: 'value'
+               value {
+                 tensor {
+                   dtype: DT_INT32
+                   tensor_shape {}
+                   int_val: 1
+                 }
+               }
+             }
+           }
+           ret { key: 'o' value: 'retval:output' })pb",
+      &fd));
+  return fd;
+}
+
+FunctionDef MakeUnaryFunction() {
+  FunctionDef fd;
+  protobuf::TextFormat::Parser parser;
+  CHECK(parser.ParseFromString(
+      R"pb(signature {
+             name: "UnaryFunction"
+             input_arg { name: "x" type: DT_INT32 }
+             output_arg { name: "ret" type: DT_INT32 }
+           }
+           node_def {
+             name: "ret"
+             op: "Identity"
+             input: "x"
+             attr {
+               key: "T"
+               value { type: DT_INT32 }
+             }
+           }
+           ret { key: "ret" value: "ret:output:0" })pb",
+      &fd));
+  return fd;
+}
+
+FunctionDef MakeBinaryFunction() {
+  FunctionDef fd;
+  protobuf::TextFormat::Parser parser;
+  CHECK(parser.ParseFromString(
+      R"pb(signature {
+             name: "BinaryFunction"
+             input_arg { name: "x" type: DT_INT32 }
+             input_arg { name: "y" type: DT_INT32 }
+             output_arg { name: "ret" type: DT_INT32 }
+           }
+           node_def {
+             name: "x_plus_y"
+             op: "AddV2"
+             input: "x"
+             input: "y"
+             attr {
+               key: "T"
+               value { type: DT_INT32 }
+             }
+           }
+           node_def {
+             name: "ret"
+             op: "Identity"
+             input: "x_plus_y:z:0"
+             attr {
+               key: "T"
+               value { type: DT_INT32 }
+             }
+           }
+           ret { key: "ret" value: "ret:output:0" })pb",
+      &fd));
+  return fd;
+}
+
+FunctionDef MakeMultiplyFunction() {
+  FunctionDef fd;
+  protobuf::TextFormat::Parser parser;
+  CHECK(parser.ParseFromString(
+      R"pb(signature {
+             name: "MultiplyFunction"
+             input_arg { name: "x" type: DT_INT32 }
+             input_arg { name: "y" type: DT_INT32 }
+             output_arg { name: "ret" type: DT_INT32 }
+           }
+           node_def {
+             name: "x_times_y"
+             op: "Mul"
+             input: "x"
+             input: "y"
+             attr {
+               key: "T"
+               value { type: DT_INT32 }
+             }
+           }
+           node_def {
+             name: "ret"
+             op: "Identity"
+             input: "x_times_y:z:0"
+             attr {
+               key: "T"
+               value { type: DT_INT32 }
+             }
+           }
+           ret { key: "ret" value: "ret:output:0" })pb",
+      &fd));
+  return fd;
+}
+
+TEST(GlobalContext, Basic) {
+  Runtime rt(GlobalEagerContext());
+  TF_ASSERT_OK(rt.CreateFunction(MakeNullaryFunction()));
+
+  absl::StatusOr<ReturnValues> rets = rt.CallFunction("NullaryFunction", {});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 1);
+}
+
+TEST(CreateTest, Call) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeNullaryFunction()));
+
+  absl::StatusOr<ReturnValues> rets = rt.CallFunction("NullaryFunction", {});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 1);
+}
+
+TEST(CreateTest, GetRoundtrip) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeNullaryFunction()));
+
+  absl::StatusOr<FunctionDef> fdef_ret = rt.GetFunctionProto("NullaryFunction");
+  TF_ASSERT_OK(fdef_ret.status());
+
+  FunctionDef fdef = *fdef_ret;
+  fdef.mutable_signature()->set_name("SecondFunction");
+
+  TF_ASSERT_OK(rt.CreateFunction(fdef));
+
+  absl::StatusOr<ReturnValues> rets = rt.CallFunction("SecondFunction", {});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 1);
+}
+
+TEST(CreateTest, MlirFromGraphDef) {
+  mlir::MLIRContext mctx;
+  mctx.getOrLoadDialect<mlir::tfg::TFGraphDialect>();
+  auto m = mlir::parseSourceString<mlir::ModuleOp>(
+      R"mlir(
+        module  {
+          tfg.func @NullaryFunction()
+               -> (tensor<i32> {tfg.dtype = i32, tfg.name = "o"})
+           {
+            %Const, %ctl = Const name("retval") {dtype = i32, value = dense<1> : tensor<i32>} : () -> (tensor<i32>)
+            return(%Const) : tensor<i32>
+          }
+        }
+      )mlir",
+      &mctx);
+
+  mlir::tfg::GraphFuncOp fop =
+      *m->getBody()->op_begin<mlir::tfg::GraphFuncOp>();
+
+  EagerContextPtr ectx = TestingEagerCtx();
+  Runtime rt(*ectx);
+  // Note: this is the price we'll have to pay until we can properly link
+  // MLIR headers into pybind wrappers (not to be confused with pybind
+  // converters, which are a separate thing - we just talk about header
+  // dependencies here).
+  OpaqueTfgGraphFuncOp* opaque_fop =
+      reinterpret_cast<OpaqueTfgGraphFuncOp*>(&fop);
+  TF_ASSERT_OK(rt.CreateFunction(opaque_fop));
+
+  absl::StatusOr<ReturnValues> rets = rt.CallFunction("NullaryFunction", {});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 1);
+}
+
+TEST(CallTest, Nullary) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeNullaryFunction()));
+
+  absl::StatusOr<ReturnValues> rets = rt.CallFunction("NullaryFunction", {});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 1);
+}
+
+TEST(CallTest, Unary) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeUnaryFunction()));
+
+  auto x = IntScalarTensor(*ctx, 1);
+  absl::StatusOr<ReturnValues> rets =
+      rt.CallFunction("UnaryFunction", {x.get()});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 1);
+}
+
+TEST(CallTest, Binary) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeBinaryFunction()));
+
+  auto x = IntScalarTensor(*ctx, 1);
+  auto y = IntScalarTensor(*ctx, 1);
+  absl::StatusOr<ReturnValues> rets =
+      rt.CallFunction("BinaryFunction", {x.get(), y.get()});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 2);
+}
+
+TEST(TransformTest, TestPassOnBinaryFunction) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeBinaryFunction()));
+
+  testing::RegisterTestPass();
+  TF_EXPECT_OK(rt.TransformFunction("BinaryFunction", "test-pass"));
+
+  auto x = IntScalarTensor(*ctx, 2);
+  auto y = IntScalarTensor(*ctx, 3);
+  absl::StatusOr<ReturnValues> rets =
+      rt.CallFunction("BinaryFunction", {x.get(), y.get()});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 6);
+}
+
+TEST(TransformTest, TestPassOnMultiplyFunction) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeMultiplyFunction()));
+
+  testing::RegisterTestPass();
+  TF_EXPECT_OK(rt.TransformFunction("MultiplyFunction", "test-pass-tf-dialect",
+                                    Runtime::Dialect::TF));
+
+  auto x = IntScalarTensor(*ctx, 2);
+  auto y = IntScalarTensor(*ctx, 3);
+  absl::StatusOr<ReturnValues> rets =
+      rt.CallFunction("MultiplyFunction", {x.get(), y.get()});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 5);
+}
+
+TEST(TransformTest, TestMixedPassesOnBinaryFunction) {
+  EagerContextPtr ctx = TestingEagerCtx();
+  Runtime rt(*ctx);
+  TF_ASSERT_OK(rt.CreateFunction(MakeBinaryFunction()));
+
+  testing::RegisterTestPass();
+  TF_EXPECT_OK(rt.TransformFunction("BinaryFunction", "test-pass"));
+  TF_EXPECT_OK(rt.TransformFunction("BinaryFunction", "test-pass-tf-dialect",
+                                    Runtime::Dialect::TF));
+
+  auto x = IntScalarTensor(*ctx, 2);
+  auto y = IntScalarTensor(*ctx, 3);
+  absl::StatusOr<ReturnValues> rets =
+      rt.CallFunction("BinaryFunction", {x.get(), y.get()});
+  TF_ASSERT_OK(rets.status());
+  ASSERT_EQ(rets->size(), 1);
+  ASSERT_EQ(rets->at(0)->DataType(), DT_INT32);
+  EXPECT_EQ(IntValue(*(rets->at(0))), 5);
+}
+
+}  // namespace
+}  // namespace function
+}  // namespace core
+}  // namespace tensorflow

tensorflow/core/function/runtime_client/runtime_client_test.py

@@ -0,0 +1,322 @@
+# 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 runtime_client."""
+
+from google.protobuf import text_format
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.function.runtime_client import runtime_client
+from tensorflow.core.function.testing import test_pass
+from tensorflow.python import tf2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import execute
+from tensorflow.python.eager import remote
+from tensorflow.python.framework import constant_op
+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 RuntimeClientTest(test.TestCase):
+
+  def test_create_nullary(self):
+    fndef = text_format.Parse(
+        """
+            signature {
+               name: 'NullaryFunction'
+               output_arg { name: 'o' type: DT_INT32 }
+             }
+             node_def {
+               name: 'retval'
+               op: 'Const'
+               attr {
+                 key: 'dtype'
+                 value { type: DT_INT32 }
+               }
+               attr {
+                 key: 'value'
+                 value {
+                   tensor {
+                     dtype: DT_INT32
+                     tensor_shape {}
+                     int_val: 1
+                   }
+                 }
+               }
+             }
+             ret { key: 'o' value: 'retval:output' }
+         """,
+        function_pb2.FunctionDef(),
+    )
+
+    ctx = runtime_client.GlobalEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    rt.CreateFunction(fndef)
+
+  def test_create_function_called_by_py_runtime(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    fndef = text_format.Parse(
+        """
+            signature {
+               name: 'NullaryFunction'
+               output_arg { name: 'o' type: DT_INT32 }
+             }
+             node_def {
+               name: 'retval'
+               op: 'Const'
+               attr {
+                 key: 'dtype'
+                 value { type: DT_INT32 }
+               }
+               attr {
+                 key: 'value'
+                 value {
+                   tensor {
+                     dtype: DT_INT32
+                     tensor_shape {}
+                     int_val: 1
+                   }
+                 }
+               }
+             }
+             ret { key: 'o' value: 'retval:output' }
+         """,
+        function_pb2.FunctionDef(),
+    )
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    rt.CreateFunction(fndef)
+
+    ret, = execute.execute("NullaryFunction", 1, [], (), context.context())
+    self.assertAllEqual(ret, 1)
+
+  def test_get_function_proto_from_py_runtime_function(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f():
+      return 1
+
+    cf = f.get_concrete_function()
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    fndef = rt.GetFunctionProto(cf.function_def.signature.name)
+
+    self.assertEqual(fndef.signature.name, cf.function_def.signature.name)
+
+  def test_concrete_function_editing_proto_executed_directly(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f():
+      return 1
+
+    cf = f.get_concrete_function()
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    fndef = rt.GetFunctionProto(cf.function_def.signature.name)
+
+    fndef.node_def[0].attr["value"].tensor.int_val[0] = 2
+
+    rt.CreateFunction(fndef)
+
+    ret, = execute.execute(fndef.signature.name, 1, [], (), context.context())
+    self.assertAllEqual(ret, 2)
+
+  def test_concrete_function_editing_proto(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f():
+      return 1
+
+    cf = f.get_concrete_function()
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    fndef = rt.GetFunctionProto(cf.function_def.signature.name)
+
+    fndef.node_def[0].attr["value"].tensor.int_val[0] = 2
+
+    rt.CreateFunction(fndef)
+
+    self.assertAllEqual(self.evaluate(f()), 2)
+
+  def test_concrete_function_editing_proto_after_instantiation(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f():
+      return 1
+
+    cf = f.get_concrete_function()
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    fndef = rt.GetFunctionProto(cf.function_def.signature.name)
+
+    fndef.node_def[0].attr["value"].tensor.int_val[0] = 2
+
+    rt.CreateFunction(fndef)
+
+    self.assertAllEqual(self.evaluate(f()), 2)
+
+  def test_concrete_function_editing_via_mlir_pass_tfg_dialect(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    one = constant_op.constant(1)
+    cf = f.get_concrete_function(one, one)
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    rt.TransformFunction(cf.function_def.signature.name, "test-pass")
+
+    # 1 + 1 = 2. But the pass changes it to 1 * 1.
+    self.assertAllEqual(self.evaluate(f(one, one)), 1)
+
+  def test_concrete_function_editing_via_mlir_pass_tf_dialect(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.multiply(x, y, name="x_times_y")
+
+    one = constant_op.constant(1)
+    cf = f.get_concrete_function(one, one)
+    fname = cf.function_def.signature.name
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+
+    # 1 * 1 = 1 -> 1 + 1 = 2
+    rt.TransformFunction(
+        fname, "test-pass-tf-dialect", runtime_client.Runtime.Dialect.TF
+    )
+
+    self.assertAllEqual(f(one, one), 2)
+
+  def test_concrete_function_editing_via_mlir_pass_mixed_dialects(self):
+    if not tf2.enabled():
+      self.skipTest("TF2 test")
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    one = constant_op.constant(1)
+    cf = f.get_concrete_function(one, one)
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    fname = cf.function_def.signature.name
+
+    # 1 + 1 = 2 -> 1 * 1 = 1
+    rt.TransformFunction(fname, "test-pass")
+
+    self.assertAllEqual(f(one, one), 1)
+
+    # 1 * 1 = 1 -> 1 + 1 = 2
+    rt.TransformFunction(
+        fname, "test-pass-tf-dialect", runtime_client.Runtime.Dialect.TF
+    )
+
+    self.assertAllEqual(f(one, one), 2)
+
+
+class RuntimeClientMultiWorkersTest(test.TestCase):
+
+  @test_util.run_v2_only
+  def setUp(self):
+    super().setUp()
+
+    workers, _ = test_util.create_local_cluster(2, 0)
+
+    remote.connect_to_remote_host(
+        [workers[0].target, workers[1].target])
+
+    self.device0 = "/job:worker/replica:0/task:0/device:CPU:0"
+    self.device1 = "/job:worker/replica:0/task:1/device:CPU:0"
+
+  @test_util.run_v2_only
+  def tearDown(self):
+    super().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()
+
+  @test_util.run_v2_only
+  def test_transform_function_in_remote_contexts(self):
+    """Tests if function_defs in remote contexts could be transformed."""
+
+    @def_function.function
+    def add(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    inputs = [1.0, 2.0]
+
+    with ops.device(self.device0):
+      result = add(*inputs)
+
+    self.assertAllEqual(result, 3.0)
+    self.assertEqual(result.device, self.device0)
+
+    with ops.device(self.device1):
+      result = add(*inputs)
+
+    self.assertAllEqual(result, 3.0)
+    self.assertEqual(result.device, self.device1)
+
+    cf = add.get_concrete_function(*inputs)
+    function_name = cf.function_def.signature.name
+
+    ctx = runtime_client.GlobalPythonEagerContext()
+    rt = runtime_client.Runtime(ctx)
+    # "test-pass" converts add to multiply.
+    rt.TransformFunction(function_name, "test-pass")
+    fndef = rt.GetFunctionProto(function_name)
+    rt.CreateFunction(fndef)
+
+    with ops.device(self.device0):
+      result = add(*inputs)
+
+    self.assertAllEqual(result, 2.0)
+    self.assertEqual(result.device, self.device0)
+
+    with ops.device(self.device1):
+      result = add(*inputs)
+
+    self.assertAllEqual(result, 2.0)
+    self.assertEqual(result.device, self.device1)
+
+
+if __name__ == "__main__":
+  context.set_soft_device_placement(False)
+  test_pass.RegisterTestPass()
+  test.main()

tensorflow/core/function/testing/BUILD

@@ -0,0 +1,55 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension")
+
+visibility = ["//tensorflow/core/function:__subpackages__"]
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = visibility,
+)
+
+licenses(["notice"])
+
+cc_library(
+    name = "test_pass_cc",
+    testonly = 1,
+    hdrs = [
+        "test_pass.h",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core/ir:Dialect",
+        "//tensorflow/core/platform:logging",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "test_pass_pybind",
+    testonly = 1,
+    srcs = ["test_pass_pybind.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "test_pass_pybind.pyi",
+    ],
+    deps = [
+        ":test_pass_cc",
+        "@pybind11",
+    ],
+)
+
+pytype_strict_library(
+    name = "test_pass_py",
+    testonly = 1,
+    srcs = [
+        "test_pass.py",
+    ],
+    visibility = visibility,
+    deps = [
+        ":test_pass_pybind",
+        "//tensorflow/python:pywrap_tensorflow",  # buildcleaner: keep (required for TF pybind)
+    ],
+)

tensorflow/core/function/testing/test_pass.h

@@ -0,0 +1,134 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FUNCTION_TESTING_TEST_PASS_H_
+#define TENSORFLOW_CORE_FUNCTION_TESTING_TEST_PASS_H_
+
+#include <memory>
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/ir/dialect.h"
+#include "tensorflow/core/ir/ops.h"
+#include "tensorflow/core/ir/tf_op_wrapper.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace core {
+namespace function {
+namespace testing {
+
+// A simple testing pass for BinaryFunction that replaces an AddV2 node named
+// `x_plus_y` with a Mul one.
+struct TestPassTfgDialect
+    : public mlir::PassWrapper<TestPassTfgDialect,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  TestPassTfgDialect() = default;
+
+  llvm::StringRef getArgument() const final { return "test-pass"; }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    mlir::OpBuilder builder(module);
+    mlir::tfg::TFGraphDialect* dialect =
+        builder.getContext()->getOrLoadDialect<mlir::tfg::TFGraphDialect>();
+
+    mlir::Operation* target = nullptr;
+    module->walk([&target](mlir::tfg::TFOp op) {
+      if (op.nameAttr() == nullptr) {
+        return;
+      }
+      if (op.name() != "x_plus_y") {
+        return;
+      }
+      target = op.getOperation();
+    });
+    DCHECK(target != nullptr);
+
+    builder.setInsertionPoint(target);
+    mlir::OperationState opstate(builder.getUnknownLoc(), "tfg.Mul");
+    opstate.operands.append(target->getOperands().begin(),
+                            target->getOperands().end());
+    opstate.types.append(target->getResultTypes().begin(),
+                         target->getResultTypes().end());
+    opstate.addAttribute("T", target->getAttr("T"));
+    opstate.addAttribute(dialect->getNameAttrIdentifier(),
+                         builder.getStringAttr("x_times_y"));
+
+    mlir::Operation* replacement = builder.create(opstate);
+    target->replaceAllUsesWith(replacement->getResults());
+    target->erase();
+  }
+};
+
+// A simple testing pass that replaces the first Mul node in the module
+// to a AddV2 node and names it `x_plus_y`.
+struct TestPassTfDialect
+    : public mlir::PassWrapper<TestPassTfDialect,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  TestPassTfDialect() = default;
+
+  llvm::StringRef getArgument() const final { return "test-pass-tf-dialect"; }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    mlir::OpBuilder builder(module);
+
+    mlir::Operation* target = nullptr;
+    module->walk([&target](mlir::Operation* op) {
+      if (op->getName().getStringRef() == "tf.Mul") {
+        target = op;
+        return;
+      }
+    });
+    DCHECK(target != nullptr);
+
+    builder.setInsertionPoint(target);
+    auto replacement = mlir::TF::AddV2Op::create(
+        builder,
+        mlir::NameLoc::get(
+            mlir::StringAttr::get(builder.getContext(), "x_plus_y")),
+        target->getResultTypes(), target->getOperand(0), target->getOperand(1));
+    target->replaceAllUsesWith(replacement->getResults());
+    target->erase();
+  }
+};
+
+inline std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfgDialectTestPass() {
+  return std::make_unique<TestPassTfgDialect>();
+}
+
+inline std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfDialectTestPass() {
+  return std::make_unique<TestPassTfDialect>();
+}
+
+inline void RegisterTestPass() {
+  mlir::registerPass([] { return CreateTfgDialectTestPass(); });
+  mlir::registerPass([] { return CreateTfDialectTestPass(); });
+}
+
+}  // namespace testing
+}  // namespace function
+}  // namespace core
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_FUNCTION_TESTING_TEST_PASS_H_

tensorflow/core/function/testing/test_pass.py

@@ -0,0 +1,21 @@
+# 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.
+# ==============================================================================
+"""Low level TF runtime client."""
+
+# TF oddity: this import loads TF-specific dynamic libraries.
+from tensorflow.python import pywrap_tensorflow  # pylint:disable=g-bad-import-order,unused-import
+from tensorflow.core.function.testing import test_pass_pybind
+
+RegisterTestPass = test_pass_pybind.RegisterTestPass

tensorflow/core/function/testing/test_pass_pybind.cc

@@ -0,0 +1,22 @@
+/* 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.
+==============================================================================*/
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/function/testing/test_pass.h"
+
+PYBIND11_MODULE(test_pass_pybind, m) {
+  m.def("RegisterTestPass",
+        &tensorflow::core::function::testing::RegisterTestPass);
+}

tensorflow/core/function/testing/test_pass_pybind.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 RegisterTestPass() -> None: ...

tensorflow/core/function/trace_type/BUILD

@@ -0,0 +1,198 @@
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "trace_type",
+    srcs = [
+        "__init__.py",
+        "trace_type_builder.py",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":custom_nest_trace_type",
+        ":default_types",
+        ":serialization",
+        ":util",
+        "//tensorflow/python/types:trace",
+        "//tensorflow/python/util:custom_nest_protocol",
+    ],
+)
+
+pytype_strict_library(
+    name = "util",
+    srcs = [
+        "util.py",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        "//third_party/py/numpy",
+    ],
+)
+
+py_test(
+    name = "trace_type_test",
+    srcs = ["trace_type_test.py"],
+    strict_deps = True,
+    # TODO(b/213375459): For continous benchmark monitoring
+    visibility = ["//learning/brain/contrib/eager/python/examples:__pkg__"],
+    deps = [
+        ":custom_nest_trace_type",
+        ":default_types",
+        ":trace_type",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:combinations",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+pytype_strict_library(
+    name = "default_types",
+    srcs = [
+        "default_types.py",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":serialization",
+        ":util",
+        "//tensorflow/core/function/trace_type:default_types_proto_py",
+        "//tensorflow/python/types:trace",
+        "//tensorflow/python/util:custom_nest_protocol",
+    ],
+)
+
+py_test(
+    name = "default_types_test",
+    srcs = ["default_types_test.py"],
+    strict_deps = True,
+    deps = [
+        ":default_types",
+        ":serialization",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/types:trace",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+pytype_strict_library(
+    name = "custom_nest_trace_type",
+    srcs = [
+        "custom_nest_trace_type.py",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = [
+        ":util",
+        "//tensorflow/python/types:trace",
+        "//tensorflow/python/util:custom_nest_protocol",
+    ],
+)
+
+py_test(
+    name = "custom_nest_trace_type_test",
+    srcs = ["custom_nest_trace_type_test.py"],
+    strict_deps = True,
+    deps = [
+        ":custom_nest_trace_type",
+        ":default_types",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/types:trace",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+pytype_strict_library(
+    name = "serialization",
+    srcs = [
+        "serialization.py",
+    ],
+    visibility = ["//tensorflow:internal"],
+    deps = ["//tensorflow/core/function/trace_type:serialization_proto_py"],
+)
+
+py_test(
+    name = "serialization_test",
+    srcs = ["serialization_test.py"],
+    strict_deps = True,
+    deps = [
+        ":serialization",
+        #internal proto upb dep
+        "//tensorflow/core/function/trace_type:serialization_test_proto_py",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tf_proto_library(
+    name = "serialization_proto",
+    srcs = [
+        "serialization.proto",
+    ],
+    visibility = ["//tensorflow:internal"],
+)
+
+tf_proto_library(
+    name = "serialization_test_proto",
+    srcs = [
+        "serialization_test.proto",
+    ],
+    protodeps = [
+        ":serialization_proto",
+    ],
+    visibility = ["//tensorflow:internal"],
+)
+
+tf_proto_library(
+    name = "default_types_proto",
+    srcs = [
+        "default_types.proto",
+    ],
+    protodeps = [
+        ":serialization_proto",
+    ],
+    visibility = ["//tensorflow:internal"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "serialization_py_pb2",
+#     visibility = ["//tensorflow:internal"],
+#     deps = [":serialization_proto"],
+# )
+#
+# py_proto_library(
+#     name = "serialization_test_py_pb2",
+#     visibility = ["//tensorflow:internal"],
+#     deps = [":serialization_test_proto"],
+# )
+#
+# py_proto_library(
+#     name = "default_types_py_pb2",
+#     visibility = ["//tensorflow:internal"],
+#     deps = [":default_types_proto"],
+# )
+# copybara:uncomment_end

tensorflow/core/function/trace_type/__init__.py

@@ -0,0 +1,38 @@
+# 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.
+# ==============================================================================
+"""Trace-time type system for tf.function (TraceType).
+
+Trace-time types describe things like tf.function signatures and type
+constraints in some ops.
+
+This module provides utilities and concrete tf.types.experimental.TraceType
+definitions for common Python types like containers, along with a generic
+implementation for Python objects.
+See also: tf.types.experimental.TraceType
+
+Other implementations of TraceType include tf.TypeSpec and its subclasses.
+"""
+
+from tensorflow.core.function.trace_type.default_types import register_tensor_type
+from tensorflow.core.function.trace_type.default_types import Weakref
+from tensorflow.core.function.trace_type.serialization import deserialize
+from tensorflow.core.function.trace_type.serialization import register_serializable
+from tensorflow.core.function.trace_type.serialization import Serializable
+from tensorflow.core.function.trace_type.serialization import serialize
+from tensorflow.core.function.trace_type.serialization import SerializedTraceType
+from tensorflow.core.function.trace_type.trace_type_builder import from_value
+from tensorflow.core.function.trace_type.trace_type_builder import InternalCastContext
+from tensorflow.core.function.trace_type.trace_type_builder import InternalPlaceholderContext
+from tensorflow.core.function.trace_type.trace_type_builder import InternalTracingContext

tensorflow/core/function/trace_type/custom_nest_trace_type.py

@@ -0,0 +1,143 @@
+# 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.
+# ==============================================================================
+"""TraceType implementations for classes thatimplement the CustomNestProtocol."""
+
+from typing import Any, Iterator, List as PythonList, Optional, Sequence, Tuple as PythonTuple, Type
+
+from tensorflow.core.function.trace_type import util
+from tensorflow.python.types import trace
+from tensorflow.python.util import custom_nest_protocol
+
+
+class CustomNestTraceType(trace.TraceType):
+  """Represents the TraceType of a class implmenting the CustomNestProtocol."""
+
+  def __init__(
+      self,
+      value_type: Type[Any],
+      metadata: Any,
+      components: PythonTuple[trace.TraceType],
+  ):
+    if not issubclass(value_type, custom_nest_protocol.CustomNestProtocol):
+      raise ValueError(f"{value_type!r} does not implement CustomNestProtocol.")
+    self.value_type = value_type
+    self.metadata = metadata
+    self.components = components
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    if not self._is_same_trace_type(other):
+      return False
+    for c_self, c_other in zip(self.components, other.components):  # pytype: disable=attribute-error
+      if not c_self.is_subtype_of(c_other):
+        return False
+    return True
+
+  def most_specific_common_supertype(
+      self, others: Sequence[trace.TraceType]
+  ) -> Optional["CustomNestTraceType"]:
+    for other in others:
+      if not self._is_same_trace_type(other):
+        return None
+
+    others_components = [other.components for other in others]  # pytype: disable=attribute-error
+    supertyped_components = tuple(
+        self_component.most_specific_common_supertype(others_component)
+        for self_component, *others_component in zip(
+            self.components, *others_components
+        )
+    )
+    return CustomNestTraceType(
+        self.value_type, self.metadata, supertyped_components
+    )
+
+  def __eq__(self, other: trace.TraceType) -> bool:
+    return (
+        isinstance(other, CustomNestTraceType)
+        and self.value_type == other.value_type
+        and self.metadata == other.metadata
+        and self.components == other.components
+    )
+
+  def __hash__(self) -> int:
+    # The hash computation doesn't use self.metadata, so unhashable metadata can
+    # be used. The `self.__eq__` method is used instead to differentiate between
+    # two objects with the same components but different metadata.
+    return hash((self.value_type, self.components))
+
+  def __repr__(self) -> str:
+    return (
+        f"{self.__class__.__name__} [metadata={self.metadata!r}, "
+        f"components={self.components!r}]"
+    )
+
+  def placeholder_value(self, placeholder_context: Any) -> Any:
+    components_placeholder_value = tuple(
+        c.placeholder_value(placeholder_context) for c in self.components
+    )
+    return self.value_type.__tf_unflatten__(
+        self.metadata, components_placeholder_value
+    )
+
+  def to_tensors(self, value: Any) -> PythonList[Any]:
+    if not isinstance(value, self.value_type):
+      raise TypeError(f"{value!r} is not of type {self.value_type}.")
+    _, value_components = value.__tf_flatten__()
+    flattened_values = []
+    for value_comp, type_comp in zip(value_components, self.components):
+      flattened_values.extend(type_comp.to_tensors(value_comp))
+    return flattened_values
+
+  def from_tensors(self, tensors: Iterator[Any]) -> Any:
+    return self.value_type.__tf_unflatten__(
+        self.metadata, tuple(c.from_tensors(tensors) for c in self.components)
+    )
+
+  def flatten(self) -> PythonList[trace.TraceType]:
+    flat_list = []
+    for c in self.components:
+      flat_list.extend(c.flatten())
+    return flat_list
+
+  def cast(self, value: Any, casting_context: Any) -> Any:
+    if not isinstance(value, self.value_type):
+      raise TypeError(f"[{value!r}] is not of type {self.value_type}.")
+    value_metadata, value_components = value.__tf_flatten__()
+    if self.metadata != value_metadata:
+      raise ValueError(
+          f"Metadata mismatch: [{self.metadata!r}] != [{value_metadata!r}]."
+      )
+    if len(self.components) != len(value_components):
+      raise ValueError(
+          f"Lengths of components mismatch: {len(self.components)} != "
+          f"{len(value_components)}."
+      )
+
+    casted_value_components, was_casted = util.cast_and_return_whether_casted(
+        self.components, value_components, casting_context
+    )
+    if was_casted:
+      return self.value_type.__tf_unflatten__(
+          self.metadata, casted_value_components
+      )
+    else:
+      return value
+
+  def _is_same_trace_type(self, other: trace.TraceType) -> bool:
+    return (
+        isinstance(other, CustomNestTraceType)
+        and self.value_type == other.value_type
+        and self.metadata == other.metadata
+        and len(self.components) == len(other.components)
+    )

tensorflow/core/function/trace_type/custom_nest_trace_type_test.py

@@ -0,0 +1,181 @@
+# 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 CustomNestTraceType."""
+
+import dataclasses
+from typing import Any
+
+from tensorflow.core.function.trace_type import custom_nest_trace_type
+from tensorflow.core.function.trace_type import default_types
+from tensorflow.python.platform import test
+from tensorflow.python.types import trace
+
+
+class MockSupertypes2With3(trace.TraceType):
+
+  def __init__(self, obj):
+    self._object = obj
+
+  def is_subtype_of(self, other):
+    return self._object == 2 and other._object == 3
+
+  def most_specific_common_supertype(self, others):
+    if not others:
+      return self
+
+    if self._object == 2 and isinstance(others[0]._object, int):
+      return MockSupertypes2With3(3)
+    else:
+      return None
+
+  def placeholder_value(self, placeholder_context=None):
+    raise NotImplementedError
+
+  def __eq__(self, other) -> bool:
+    return isinstance(other, type(self)) and self._object == other._object
+
+  def __hash__(self) -> int:
+    return self._object_hash
+
+
+class Mock2AsTopType(MockSupertypes2With3):
+
+  def is_subtype_of(self, other):
+    return other._object == 2
+
+  def most_specific_common_supertype(self, others):
+    if not all(isinstance(other, Mock2AsTopType) for other in others):
+      return None
+    return (
+        self
+        if all(self._object == other._object for other in others)
+        else Mock2AsTopType(2)
+    )
+
+
+@dataclasses.dataclass
+class MaskedValuePair:
+  mask: bool
+  value1: Any
+  value2: Any
+
+  def __tf_flatten__(self):
+    metadata = (self.mask,)
+    components = (self.value1, self.value2)
+    return metadata, components
+
+  @classmethod
+  def __tf_unflatten__(cls, metadata, leaves):
+    mask = metadata[0]
+    value1 = leaves[0]
+    value2 = leaves[1]
+    return MaskedValuePair(mask=mask, value1=value1, value2=value2)
+
+
+class CustomNestTraceTypeTest(test.TestCase):
+
+  def testCustomNestTraceTypeEq(self):
+    trace_components1 = (default_types.Literal(1.0), default_types.Literal(2.0))
+    t1 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components1
+    )
+
+    trace_components2 = (default_types.Literal(1.0), default_types.Literal(2.0))
+    t2 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components2
+    )
+    self.assertEqual(t1, t2)
+
+    trace_components3 = (default_types.Literal(2), default_types.Literal(1))
+    t3 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (False,), trace_components3
+    )
+    self.assertNotEqual(t1, t3)
+
+  def testCustomNestTraceTypePlaceholderValue(self):
+    trace_components1 = (default_types.Literal(1.0), default_types.Literal(2.0))
+    t1 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components1
+    )
+
+    mvp1 = MaskedValuePair(mask=True, value1=1.0, value2=2.0)
+    phv1 = t1.placeholder_value(None)
+    self.assertEqual(phv1.mask, mvp1.mask)
+    self.assertEqual(phv1.value1, mvp1.value1)
+    self.assertEqual(phv1.value2, mvp1.value2)
+
+    trace_components2 = (default_types.Literal(2), default_types.Literal(1))
+    t2 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (False,), trace_components2
+    )
+    mvp2 = MaskedValuePair(mask=False, value1=2, value2=1)
+    phv2 = t2.placeholder_value(None)
+    self.assertEqual(phv2.mask, mvp2.mask)
+    self.assertEqual(phv2.value1, mvp2.value1)
+    self.assertEqual(phv2.value2, mvp2.value2)
+
+  def testCustomNestTraceTypeSubtype(self):
+    trace_components1 = (Mock2AsTopType(1), Mock2AsTopType(1))
+    t1 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components1
+    )
+
+    trace_components2 = (Mock2AsTopType(2), Mock2AsTopType(2))
+    t2 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components2
+    )
+    self.assertTrue(t1.is_subtype_of(t2))
+    self.assertFalse(t2.is_subtype_of(t1))
+
+    trace_components3 = (Mock2AsTopType(2), Mock2AsTopType(2))
+    t4 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components3
+    )
+    self.assertTrue(t2.is_subtype_of(t4))
+    self.assertTrue(t4.is_subtype_of(t2))
+
+    trace_components4 = (Mock2AsTopType(1), Mock2AsTopType(2))
+    t4 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components4
+    )
+    self.assertFalse(t1.is_subtype_of(t4))
+    self.assertFalse(t4.is_subtype_of(t1))
+
+  def testCustomNestTraceTypeSupertype(self):
+    trace_components1 = (MockSupertypes2With3(2), MockSupertypes2With3(2))
+    t1 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components1
+    )
+
+    trace_components2 = (MockSupertypes2With3(1), MockSupertypes2With3(1))
+    t2 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components2
+    )
+
+    trace_components3 = (MockSupertypes2With3(3), MockSupertypes2With3(3))
+    t3 = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), trace_components3
+    )
+    self.assertEqual(t1.most_specific_common_supertype([t2]), t3)
+
+    c_super_none = custom_nest_trace_type.CustomNestTraceType(
+        MaskedValuePair, (True,), (None, None)
+    )
+    self.assertEqual(t2.most_specific_common_supertype([t3]), c_super_none)
+    self.assertEqual(t3.most_specific_common_supertype([t1]), c_super_none)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/core/function/trace_type/default_types.proto

@@ -0,0 +1,63 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto2";
+
+package tensorflow.core.function.trace_type.default_types;
+
+import "tensorflow/core/function/trace_type/serialization.proto";
+
+// Represents a serialized Literal type.
+message SerializedLiteral {
+  message NoneValue {}
+
+  oneof value {
+    bool bool_value = 1;
+    int64 int_value = 2;
+    double float_value = 3;  // Python "float" is double precision.
+    string str_value = 4;
+    NoneValue none_value = 5;
+  }
+}
+
+// Represents a serialized Tuple type.
+message SerializedTuple {
+  repeated tensorflow.core.function.trace_type.serialization.SerializedTraceType
+      components = 1;
+}
+
+// Represents a serialized List type.
+message SerializedList {
+  optional SerializedTuple components_tuple = 1;
+}
+
+// Represents a serialized NamedTuple type.
+message SerializedNamedTuple {
+  optional string type_name = 1;
+  repeated string attribute_names = 2;
+  optional SerializedTuple attributes = 3;
+}
+
+// Represents a serialized Attrs type.
+message SerializedAttrs {
+  optional SerializedNamedTuple named_attributes = 1;
+}
+
+// Represents a serialized Dict type.
+message SerializedDict {
+  repeated SerializedLiteral keys = 1;
+  repeated tensorflow.core.function.trace_type.serialization.SerializedTraceType
+      values = 2;
+}

tensorflow/core/function/trace_type/default_types.py

@@ -0,0 +1,826 @@
+# 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.
+# ==============================================================================
+"""TraceType implementations for common Python types."""
+
+import collections
+import math
+import numbers
+from typing import Any, Dict as PythonDict, Hashable, List as PythonList, Optional, Sequence, Tuple as PythonTuple, Type
+import weakref
+
+from tensorflow.core.function.trace_type import default_types_pb2
+from tensorflow.core.function.trace_type import serialization
+from tensorflow.core.function.trace_type import util
+from tensorflow.python.types import trace
+
+# Register the TraceType of Tensor (aka TensorSpec) to avoid cyclic dependency.
+TENSOR = None
+
+
+def register_tensor_type(tensor_type):
+  global TENSOR
+  if not TENSOR:
+    TENSOR = tensor_type
+  else:
+    raise AssertionError("Tensor type is already registered.")
+
+NanMarker = object()
+
+
+def is_nan(x):
+  """Checks if given value is a Python NaN."""
+  if not isinstance(x, numbers.Number):
+    return False
+
+  if isinstance(x, complex):
+    return math.isnan(x.real) or math.isnan(x.imag)
+  else:
+    return math.isnan(x)
+
+
+class Literal(trace.TraceType, serialization.Serializable):
+  """Represents a Literal type like bool, int or string."""
+
+  def __init__(self, value: Any):
+    # We match nan values against each other even though Python doesn't.
+    if is_nan(value):
+      value = NanMarker
+
+    self.value = value
+    self._value_hash = hash(value)
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    return self == other
+
+  def most_specific_common_supertype(
+      self, types: Sequence[trace.TraceType]) -> Optional["Literal"]:
+    return self if all(self == other for other in types) else None
+
+  @classmethod
+  def experimental_type_proto(cls) -> Type[default_types_pb2.SerializedLiteral]:
+    return default_types_pb2.SerializedLiteral
+
+  @classmethod
+  def experimental_from_proto(
+      cls, proto: default_types_pb2.SerializedLiteral) -> "Literal":
+    if proto.HasField("bool_value"):
+      return Literal(proto.bool_value)
+
+    if proto.HasField("int_value"):
+      return Literal(proto.int_value)
+
+    if proto.HasField("float_value"):
+      return Literal(proto.float_value)
+
+    if proto.HasField("str_value"):
+      return Literal(proto.str_value)
+
+    if proto.HasField("none_value"):
+      return Literal(None)
+
+    raise ValueError("Malformed Literal proto can not be deserialized")
+
+  def experimental_as_proto(self) -> default_types_pb2.SerializedLiteral:
+    if isinstance(self.value, bool):
+      return default_types_pb2.SerializedLiteral(bool_value=self.value)
+
+    if isinstance(self.value, int):
+      return default_types_pb2.SerializedLiteral(int_value=self.value)
+
+    if isinstance(self.value, float):
+      return default_types_pb2.SerializedLiteral(float_value=self.value)
+
+    if isinstance(self.value, str):
+      return default_types_pb2.SerializedLiteral(str_value=self.value)
+
+    if self.value is None:
+      return default_types_pb2.SerializedLiteral(
+          none_value=default_types_pb2.SerializedLiteral.NoneValue())
+
+    raise ValueError("Can not serialize Literal of type " +
+                     type(self.value).__name__)
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    # TODO(b/263505796): Remove this check when a range's placeholder output
+    # is expected to be a range and not a list.
+    if isinstance(self.value, range):
+      return list(self.value)
+
+    if self.value is NanMarker:
+      return float("nan")
+
+    return self.value
+
+  def cast(self, value: Any, casting_context: Any) -> Any:
+    if self.value is NanMarker and is_nan(value):
+      return value
+
+    if value == self.value:
+      return value
+    else:
+      raise ValueError(f"Can not cast {value!r} to {self!r}")
+
+  def __eq__(self, other) -> bool:
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    return isinstance(other, Literal) and self.value == other.value
+
+  def __hash__(self) -> int:
+    return self._value_hash
+
+  def __repr__(self) -> str:
+    return f"{self.__class__.__name__}[{self.value!r}]"
+
+
+class Weakref(trace.TraceType):
+  """Represents weakref of an arbitrary Python object.
+
+  When a function argument is a custom class, instead of making a copy of it
+  just for the sake of function cache, a weakref is instead kept to save memory.
+  """
+
+  def __init__(self, ref: weakref.ReferenceType):
+    self._ref = ref
+    self._ref_hash = hash(ref)
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    return self == other
+
+  def most_specific_common_supertype(
+      self, types: Sequence[trace.TraceType]) -> Optional["Weakref"]:
+    return self if all(self == other for other in types) else None
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    return self._ref()
+
+  def cast(self, value, _):
+    if value is self._ref() or value == self._ref():
+      return value
+
+    # We unwrap objects when generating the TraceType so we allow matching now.
+    while hasattr(value, "__wrapped__"):
+      value = value.__wrapped__
+      if value is self._ref():
+        return value
+
+    raise ValueError(f"Can not cast {value!r} to {self!r}")
+
+  def __eq__(self, other):
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    if not isinstance(other, Weakref):
+      return False
+
+    if self._ref() is None or other._ref() is None:
+      return False
+
+    if self._ref() is other._ref():
+      return True
+
+    return self._ref == other._ref
+
+  def __hash__(self):
+    return self._ref_hash
+
+  def __repr__(self) -> str:
+    return f"{self.__class__.__name__}[{self._ref!r}])"
+
+
+class Tuple(trace.TraceType, serialization.Serializable):
+  """Represents a tuple of TraceType objects."""
+
+  def __init__(self, *components: trace.TraceType):
+    self.components = components
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    if (not isinstance(other, Tuple) or
+        len(self.components) != len(other.components)):
+      return False
+
+    return all(
+        self_component.is_subtype_of(other_component) for self_component,
+        other_component in zip(self.components, other.components))
+
+  def most_specific_common_supertype(
+      self, others: Sequence[trace.TraceType]) -> Optional["Tuple"]:
+    """See base class."""
+    if not all(
+        isinstance(other, Tuple) and
+        len(self.components) == len(other.components) for other in others):
+      return None
+
+    supertyped_components = []
+    for i, component in enumerate(self.components):
+      supertyped_component = component.most_specific_common_supertype(
+          [other.components[i] for other in others])
+      if supertyped_component is None:
+        return None
+      supertyped_components.append(supertyped_component)
+
+    return Tuple(*supertyped_components)
+
+  @classmethod
+  def experimental_type_proto(cls) -> Type[default_types_pb2.SerializedTuple]:
+    return default_types_pb2.SerializedTuple
+
+  @classmethod
+  def experimental_from_proto(
+      cls, proto: default_types_pb2.SerializedTuple) -> "Tuple":
+    return Tuple(*[serialization.deserialize(c) for c in proto.components])
+
+  def experimental_as_proto(self) -> default_types_pb2.SerializedTuple:
+    return default_types_pb2.SerializedTuple(
+        components=[serialization.serialize(c) for c in self.components])
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    components = [
+        component.placeholder_value(placeholder_context)
+        for component in self.components
+    ]
+    return tuple(components)
+
+  def to_tensors(self, value) -> Any:
+    assert isinstance(value, tuple)
+    flattened_values = []
+    for comp_value, comp_type in zip(value, self.components):
+      flattened_values.extend(comp_type.to_tensors(comp_value))
+    return flattened_values
+
+  def from_tensors(self, tensors) -> Any:
+    return tuple(c.from_tensors(tensors) for c in self.components)
+
+  def flatten(self) -> PythonList[trace.TraceType]:
+    flattened_types = []
+    for component in self.components:
+      flattened_types.extend(component.flatten())
+    return flattened_types
+
+  def cast(self, value: Any, casting_context) -> Any:
+    assert isinstance(value, tuple), f"Can not cast {value!r} to tuple type."
+    assert len(value) == len(
+        self.components
+    ), f"Expected {value} to have length of {len(self.components)}"
+
+    casted_values, was_casted = util.cast_and_return_whether_casted(
+        self.components, value, casting_context
+    )
+    if was_casted:
+      return tuple(casted_values)
+    else:
+      return value
+
+  def __eq__(self, other: Any) -> bool:
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    if not isinstance(other, Tuple):
+      return False
+
+    return self.components == other.components
+
+  def __hash__(self) -> int:
+    return hash(self.components)
+
+  def __repr__(self) -> str:
+    return f"Tuple[{', '.join(map(repr, self.components))}]"
+
+
+class List(trace.TraceType, serialization.Serializable):
+  """Represents a list of TraceType objects."""
+
+  def __init__(self, *components: trace.TraceType):
+    self.components_tuple = Tuple(*components)
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    if not isinstance(other, List):
+      return False
+
+    return self.components_tuple.is_subtype_of(other.components_tuple)
+
+  def most_specific_common_supertype(
+      self, others: Sequence[trace.TraceType]) -> Optional["Tuple"]:
+    """See base class."""
+    if not all(isinstance(other, List) for other in others):
+      return None
+
+    supertyped_components_tuple = (
+        self.components_tuple.most_specific_common_supertype(
+            [other.components_tuple for other in others]
+        )
+    )
+
+    if supertyped_components_tuple is None:
+      return None
+
+    return List(*supertyped_components_tuple.components)
+
+  @classmethod
+  def experimental_type_proto(cls) -> Type[default_types_pb2.SerializedList]:
+    return default_types_pb2.SerializedList
+
+  @classmethod
+  def experimental_from_proto(
+      cls, proto: default_types_pb2.SerializedList) -> "List":
+    return List(
+        *Tuple.experimental_from_proto(proto.components_tuple).components)
+
+  def experimental_as_proto(self) -> default_types_pb2.SerializedList:
+    return default_types_pb2.SerializedList(
+        components_tuple=self.components_tuple.experimental_as_proto())
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    return list(self.components_tuple.placeholder_value(placeholder_context))
+
+  def to_tensors(self, value):
+    assert isinstance(value, list)
+    return self.components_tuple.to_tensors(tuple(value))
+
+  def from_tensors(self, tensors) -> Any:
+    return list(self.components_tuple.from_tensors(tensors))
+
+  def flatten(self) -> PythonList[trace.TraceType]:
+    return self.components_tuple.flatten()
+
+  def cast(self, value: Any, casting_context) -> Any:
+    assert isinstance(value, list), f"Can not cast {value!r} to list type."
+    assert len(value) == len(
+        self.components_tuple.components
+    ), f"Expected {value} to have length of {len(self.components_tuple)}"
+
+    casted_values, was_casted = util.cast_and_return_whether_casted(
+        self.components_tuple.components, value, casting_context
+    )
+    if was_casted:
+      return list(casted_values)
+    else:
+      return value
+
+  def __eq__(self, other: Any) -> bool:
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    if not isinstance(other, List):
+      return False
+
+    return self.components_tuple == other.components_tuple
+
+  def __hash__(self) -> int:
+    return hash(self.components_tuple)
+
+  def __repr__(self) -> str:
+    return f"List[{', '.join(map(repr, self.components_tuple.components))}]"
+
+
+class NamedTuple(trace.TraceType, serialization.Serializable):
+  """Represents a NamedTuple of TraceType objects."""
+
+  def __init__(self,
+               type_name: str,
+               attribute_names: PythonTuple[str],
+               attributes: PythonTuple[trace.TraceType],
+               placeholder_type: Optional[Type[Any]] = None):
+    self.type_name = type_name
+    self.attribute_names = attribute_names
+    self.attributes = Tuple(*attributes)
+    self._placeholder_type = placeholder_type
+
+  @classmethod
+  def from_type_and_attributes(
+      cls, named_tuple_type: Any,
+      attributes: PythonTuple[trace.TraceType]) -> "NamedTuple":
+    return NamedTuple(named_tuple_type.__name__, named_tuple_type._fields,
+                      attributes, named_tuple_type)
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    if not isinstance(other, NamedTuple):
+      return False
+
+    return (self.type_name == other.type_name and
+            self.attribute_names == other.attribute_names and
+            self.attributes.is_subtype_of(other.attributes))
+
+  def most_specific_common_supertype(
+      self, others: Sequence[trace.TraceType]) -> Optional["NamedTuple"]:
+    """See base class."""
+    if not all(
+        isinstance(other, NamedTuple) and self.type_name == other.type_name and
+        self.attribute_names == other.attribute_names for other in others):
+      return None
+
+    supertyped_attributes = self.attributes.most_specific_common_supertype(
+        [other.attributes for other in others])
+
+    if supertyped_attributes is None:
+      return None
+
+    return NamedTuple(self.type_name, self.attribute_names,
+                      supertyped_attributes.components, self._placeholder_type)
+
+  @classmethod
+  def experimental_type_proto(
+      cls) -> Type[default_types_pb2.SerializedNamedTuple]:
+    return default_types_pb2.SerializedNamedTuple
+
+  @classmethod
+  def experimental_from_proto(
+      cls, proto: default_types_pb2.SerializedNamedTuple) -> "NamedTuple":
+    return NamedTuple(
+        proto.type_name, tuple(proto.attribute_names),
+        Tuple.experimental_from_proto(proto.attributes).components)
+
+  def experimental_as_proto(self) -> default_types_pb2.SerializedNamedTuple:
+    return default_types_pb2.SerializedNamedTuple(
+        type_name=self.type_name,
+        attribute_names=list(self.attribute_names),
+        attributes=self.attributes.experimental_as_proto())
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    if self._placeholder_type is None:
+      # We don't need to trace after serialization so it is not needed but we
+      # can generate a placeholder type using the description if ever needed.
+      raise ValueError("Can not generate placeholder value for NamedTuple with"
+                       " unspecified placeholder_type. Note: placeholder_type "
+                       "is lost during serialization.")
+    attribute_placeholders = [
+        attribute.placeholder_value(placeholder_context)
+        for attribute in self.attributes.components
+    ]
+    return self._placeholder_type(*attribute_placeholders)
+
+  def to_tensors(self, value: Any):
+    assert util.is_namedtuple(value)
+    flattened_values = []
+    for attribute_name, attribute_type in zip(
+        self.attribute_names, self.attributes.components):
+      attribute_value = getattr(value, attribute_name)
+      flattened_values.extend(attribute_type.to_tensors(attribute_value))
+    return flattened_values
+
+  def from_tensors(self, tensors) -> Any:
+    if self._placeholder_type is None:
+      raise ValueError("Packing serialized NamedTuples is not supported.")
+
+    return self._placeholder_type(
+        *[c.from_tensors(tensors) for c in self.attributes.components]
+    )
+
+  def flatten(self) -> PythonList[trace.TraceType]:
+    flattened_types = []
+
+    for component in self.attributes.components:
+      flattened_types.extend(component.flatten())
+
+    return flattened_types
+
+  def cast(self, value: Any, casting_context) -> Any:
+    # Value must have same attributes with the TraceType
+    assert util.is_namedtuple(
+        value
+    ), f"Cannot cast {value!r} to type {self._placeholder_type!r}."
+    value_dict = value._asdict()
+    assert set(value_dict.keys()) == set(
+        self.attribute_names
+    ), f"{value!r} has different attributes with the TraceType {self!r}"
+
+    casted_values, was_casted = util.cast_and_return_whether_casted(
+        self.attributes.components,
+        [getattr(value, name) for name in self.attribute_names],
+        casting_context,
+    )
+    if was_casted:
+      return self._placeholder_type(*casted_values)
+    else:
+      return value
+
+  def __hash__(self) -> int:
+    return hash((self.type_name, self.attribute_names, self.attributes))
+
+  def __eq__(self, other: Any) -> bool:
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    if not isinstance(other, NamedTuple):
+      return False
+
+    return (self.type_name == other.type_name and
+            self.attribute_names == other.attribute_names and
+            self.attributes == other.attributes)
+
+  def __repr__(self) -> str:
+    paired = [
+        f"[{n!r}, {c!r}]"
+        for n, c in zip(self.attribute_names, self.attributes.components)
+    ]
+    return f"{self.type_name}[{', '.join(paired)}]"
+
+
+class Attrs(trace.TraceType):
+  """Represents a class annotated by attr.s."""
+
+  def __init__(self,
+               type_name: str,
+               attribute_names: PythonTuple[str],
+               attributes: PythonTuple[trace.TraceType],
+               placeholder_type: Optional[Type[Any]] = None):
+    self.named_attributes = NamedTuple(type_name, attribute_names, attributes)
+    self._placeholder_type = placeholder_type
+
+  @classmethod
+  def from_type_and_attributes(
+      cls, attrs_type: Any,
+      attributes: PythonTuple[trace.TraceType]) -> "Attrs":
+    return Attrs(attrs_type.__name__,
+                 tuple(attr.name for attr in attrs_type.__attrs_attrs__),
+                 attributes, attrs_type)
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    if not isinstance(other, Attrs):
+      return False
+
+    return self.named_attributes.is_subtype_of(other.named_attributes)
+
+  def most_specific_common_supertype(
+      self, others: Sequence[trace.TraceType]) -> Optional["Attrs"]:
+    """See base class."""
+    if not all(isinstance(other, Attrs) for other in others):
+      return None
+
+    supertyped_attributes = (
+        self.named_attributes.most_specific_common_supertype(
+            [other.named_attributes for other in others]
+        )
+    )
+
+    if supertyped_attributes is None:
+      return None
+
+    return Attrs(self.named_attributes.type_name,
+                 self.named_attributes.attribute_names,
+                 supertyped_attributes.attributes.components,
+                 self._placeholder_type)
+
+  @classmethod
+  def experimental_type_proto(cls) -> Type[default_types_pb2.SerializedAttrs]:
+    return default_types_pb2.SerializedAttrs
+
+  @classmethod
+  def experimental_from_proto(
+      cls, proto: default_types_pb2.SerializedAttrs) -> "Attrs":
+    return Attrs(
+        proto.named_attributes.type_name,
+        tuple(proto.named_attributes.attribute_names),
+        Tuple.experimental_from_proto(
+            proto.named_attributes.attributes).components)
+
+  def experimental_as_proto(self) -> default_types_pb2.SerializedAttrs:
+    return default_types_pb2.SerializedAttrs(
+        named_attributes=self.named_attributes.experimental_as_proto())
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    if self._placeholder_type is None:
+      # We don't need to trace after serialization so it is not needed but we
+      # can generate a placeholder type using the description if ever needed.
+      raise ValueError("Can not generate placeholder value for Attrs with"
+                       " unspecified placeholder_type. Note: placeholder_type "
+                       "is lost during serialization.")
+    attribute_placeholders = [
+        attribute.placeholder_value(placeholder_context)
+        for attribute in self.named_attributes.attributes.components
+    ]
+    return self._placeholder_type(*attribute_placeholders)
+
+  def to_tensors(self, value: Any):
+    assert util.is_attrs(value)
+    flattened_values = []
+    for attribute_name, attribute_type in zip(
+        self.named_attributes.attribute_names,
+        self.named_attributes.attributes.components):
+      attribute_value = getattr(value, attribute_name)
+      flattened_values.extend(attribute_type.to_tensors(attribute_value))
+    return flattened_values
+
+  def from_tensors(self, tensors):
+    if self._placeholder_type is None:
+      raise ValueError("Packing serialized NamedTuples is not supported.")
+
+    return self._placeholder_type(
+        *[
+            c.from_tensors(tensors)
+            for c in self.named_attributes.attributes.components
+        ]
+    )
+
+  def flatten(self) -> PythonList[trace.TraceType]:
+    flattened_types = []
+
+    for component in self.named_attributes.attributes.components:
+      flattened_types.extend(component.flatten())
+
+    return flattened_types
+
+  def cast(self, value: Any, casting_context) -> Any:
+    assert util.is_attrs(value)
+
+    attr_names = self.named_attributes.attribute_names
+    casted_values, was_casted = util.cast_and_return_whether_casted(
+        self.named_attributes.attributes.components,
+        [getattr(value, name) for name in attr_names],
+        casting_context,
+    )
+
+    if was_casted:
+      return self._placeholder_type(*casted_values)
+    else:
+      return value
+
+  def __hash__(self) -> int:
+    return hash(self.named_attributes)
+
+  def __eq__(self, other: Any) -> bool:
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    if not isinstance(other, Attrs):
+      return False
+
+    return self.named_attributes == other.named_attributes
+
+  def __repr__(self) -> str:
+    name_component_zip = zip(
+        self.named_attributes.attribute_names,
+        self.named_attributes.attributes.components,
+    )
+    paired = [f"[{n!r}, {c!r}]" for n, c in name_component_zip]
+    return f"{self.named_attributes.type_name}[{', '.join(paired)}]"
+
+
+class Dict(trace.TraceType, serialization.Serializable):
+  """Represents a dictionary of TraceType objects.
+
+  Attributes:
+    mapping: A mapping from keys to corresponding TraceTypes of the dict values.
+  """
+
+  def __init__(self,
+               mapping: PythonDict[Hashable, trace.TraceType],
+               placeholder_type: Optional[Type[Any]] = None):
+    self.mapping = mapping
+    self._placeholder_type = placeholder_type
+
+  def _has_same_structure(self, other):
+    if not isinstance(other, Dict):
+      return False
+
+    return self.mapping.keys() == other.mapping.keys()
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    """See base class."""
+    if not self._has_same_structure(other):
+      return False
+
+    # We need all keys to be present because there can be logic relying on
+    # their existence or lack thereof and hence can not guarantee subtype based
+    # on a subset or superset of keys.
+    # Only the tracing code can explicitly check for key dependencies and inform
+    # that decision.
+    return all(self.mapping[key].is_subtype_of(other.mapping[key])
+               for key in self.mapping)
+
+  def most_specific_common_supertype(
+      self, types: Sequence[trace.TraceType]) -> Optional["Dict"]:
+    """See base class."""
+    if not all(self._has_same_structure(other) for other in types):
+      return None
+
+    new_mapping = {}
+    for key in self.mapping.keys():
+      common = self.mapping[key].most_specific_common_supertype(
+          [other.mapping[key] for other in types])
+      if common is None:
+        return None
+      else:
+        new_mapping[key] = common
+
+    return Dict(new_mapping, self._placeholder_type)
+
+  @classmethod
+  def experimental_type_proto(cls) -> Type[default_types_pb2.SerializedDict]:
+    return default_types_pb2.SerializedDict
+
+  @classmethod
+  def experimental_from_proto(
+      cls, proto: default_types_pb2.SerializedDict) -> "Dict":
+    return Dict({
+        Literal.experimental_from_proto(k).value: serialization.deserialize(v)
+        for k, v in zip(proto.keys, proto.values)
+    })
+
+  def experimental_as_proto(self) -> default_types_pb2.SerializedDict:
+    return default_types_pb2.SerializedDict(
+        keys=[Literal(k).experimental_as_proto() for k in self.mapping.keys()],
+        values=[serialization.serialize(v) for v in self.mapping.values()])
+
+  def placeholder_value(self, placeholder_context) -> Any:
+    if self._placeholder_type is None:
+      raise ValueError("Can not generate placeholder value for Dict with"
+                       " unspecified placeholder_type. Note: placeholder_type "
+                       "is lost during serialization.")
+    attribute_placeholders = [
+        (key, value.placeholder_value(placeholder_context))
+        for key, value in self.mapping.items()
+    ]
+    if self._placeholder_type is collections.defaultdict:
+      return dict(attribute_placeholders)
+    return self._placeholder_type(attribute_placeholders)
+
+  def to_tensors(self, value: Any):
+    assert isinstance(value, collections.abc.Mapping)
+    flattened_values = []
+    for key in sorted(self.mapping.keys()):
+      comp_value, comp_type = value[key], self.mapping[key]
+      flattened_values.extend(comp_type.to_tensors(comp_value))
+    return flattened_values
+
+  def from_tensors(self, tensors):
+    if self._placeholder_type is None:
+      raise ValueError("Packing serialized Dict is not supported.")
+
+    sorted_traversal = {
+        key: self.mapping[key].from_tensors(tensors)
+        for key in sorted(self.mapping)
+    }
+
+    if self._placeholder_type is collections.defaultdict:
+      return {key: sorted_traversal[key] for key in self.mapping}
+
+    return self._placeholder_type(
+        (key, sorted_traversal[key]) for key in self.mapping
+    )
+
+  def flatten(self) -> PythonList[trace.TraceType]:
+    flattened_types = []
+
+    for key in sorted(self.mapping.keys()):
+      flattened_types.extend(self.mapping[key].flatten())
+
+    return flattened_types
+
+  def cast(self, value: Any, casting_context) -> Any:
+    # Value must have same keys with the TraceType
+    assert isinstance(
+        value, collections.abc.Mapping
+    ), f"Can not cast {value!r} to a Dict type."
+    assert set(value.keys()) == set(
+        self.mapping.keys()
+    ), f"{value!r} has different keys with the TraceType {self!r}."
+
+    casted_values, was_casted = util.cast_and_return_whether_casted(
+        self.mapping.values(),
+        [value[k] for k in self.mapping.keys()],
+        casting_context,
+    )
+
+    if was_casted:
+      return self._placeholder_type(
+          **{k: v for k, v in zip(self.mapping.keys(), casted_values)}
+      )
+    else:
+      return value
+
+  def __eq__(self, other) -> bool:
+    if not isinstance(other, trace.TraceType):
+      return NotImplemented
+
+    if not isinstance(other, Dict):
+      return False
+
+    return self.mapping == other.mapping
+
+  def __hash__(self) -> int:
+    return hash(frozenset(self.mapping.keys()))
+
+  def __repr__(self) -> str:
+    paired = [f"[{n!r}, {t!r}]" for n, t in self.mapping.items()]
+    return f"{self.__class__.__name__}[{', '.join(paired)}]"
+
+
+serialization.register_serializable(Literal)
+serialization.register_serializable(Tuple)
+serialization.register_serializable(List)
+serialization.register_serializable(NamedTuple)
+serialization.register_serializable(Attrs)
+serialization.register_serializable(Dict)

tensorflow/core/function/trace_type/default_types_test.py

@@ -0,0 +1,396 @@
+# 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 default_types."""
+
+import collections
+
+from tensorflow.core.function.trace_type import default_types
+from tensorflow.core.function.trace_type import serialization
+from tensorflow.python.platform import test
+from tensorflow.python.types import trace
+
+
+class MockSupertypes2With3(trace.TraceType):
+
+  def __init__(self, obj):
+    self._object = obj
+
+  def is_subtype_of(self, other):
+    return self._object == 2 and other._object == 3
+
+  def most_specific_common_supertype(self, others):
+    if not others:
+      return self
+
+    if self._object == 2 and isinstance(others[0]._object, int):
+      return MockSupertypes2With3(3)
+    else:
+      return None
+
+  def placeholder_value(self, placeholder_context=None):
+    raise NotImplementedError
+
+  def __eq__(self, other) -> bool:
+    return isinstance(other, type(self)) and self._object == other._object
+
+  def __hash__(self) -> int:
+    return self._object_hash
+
+  def __repr__(self) -> str:
+    return 'MockSupertypes2With3'
+
+
+class Mock2AsTopType(MockSupertypes2With3):
+
+  def is_subtype_of(self, other):
+    return other._object == 2
+
+  def most_specific_common_supertype(self, others):
+    if not all(isinstance(other, Mock2AsTopType) for other in others):
+      return None
+    return (
+        self
+        if all(self._object == other._object for other in others)
+        else Mock2AsTopType(2)
+    )
+
+  def __repr__(self) -> str:
+    return 'Mock2AsTopType'
+
+
+class TestAttr:
+  """Helps test attrs collections."""
+
+  def __init__(self, name):
+    self.name = name
+
+
+class TestAttrsClass:
+  """Helps test attrs collections."""
+
+  __attrs_attrs__ = (TestAttr('a'), TestAttr('b'))
+
+  def __init__(self, a, b):
+    self.a = a
+    self.b = b
+
+  def __eq__(self, other):
+    return (
+        isinstance(other, TestAttrsClass)
+        and self.a == other.a
+        and self.b == other.b
+    )
+
+
+class DefaultTypesTest(test.TestCase):
+
+  def testLiteralNan(self):
+    nan_literal = default_types.Literal(float('nan'))
+    complex_nan = default_types.Literal(complex(float('nan'), 1))
+    complex_nan_other = default_types.Literal(complex(1, float('nan')))
+    self.assertEqual(nan_literal, nan_literal)
+    self.assertEqual(nan_literal, complex_nan)
+    self.assertEqual(nan_literal, complex_nan_other)
+
+  def testLiteralSupertypes(self):
+    literal_a = default_types.Literal(1)
+    literal_b = default_types.Literal(2)
+    literal_c = default_types.Literal(1)
+
+    self.assertEqual(literal_a, literal_a.most_specific_common_supertype([]))
+    self.assertEqual(
+        literal_a, literal_a.most_specific_common_supertype([literal_a])
+    )
+    self.assertEqual(
+        literal_a, literal_a.most_specific_common_supertype([literal_c])
+    )
+    self.assertIsNone(literal_a.most_specific_common_supertype([literal_b]))
+
+  def testLiteralSerialization(self):
+    literal_bool = default_types.Literal(True)
+    literal_int = default_types.Literal(1)
+    literal_float = default_types.Literal(1.2)
+    literal_str = default_types.Literal('a')
+    literal_none = default_types.Literal(None)
+
+    self.assertEqual(str(literal_bool), 'Literal[True]')
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(literal_bool)),
+        literal_bool,
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(literal_int)),
+        literal_int,
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(literal_float)),
+        literal_float,
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(literal_str)),
+        literal_str,
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(literal_none)),
+        literal_none,
+    )
+
+  def testListSupertype(self):
+    list_a = default_types.List(
+        MockSupertypes2With3(1),
+        MockSupertypes2With3(2),
+        MockSupertypes2With3(3),
+    )
+    list_b = default_types.List(
+        MockSupertypes2With3(2),
+        MockSupertypes2With3(2),
+        MockSupertypes2With3(2),
+    )
+
+    self.assertEqual(list_a, list_a.most_specific_common_supertype([]))
+    self.assertIsNone(list_a.most_specific_common_supertype([list_b]))
+    self.assertEqual(
+        list_b.most_specific_common_supertype([list_a]),
+        default_types.List(
+            MockSupertypes2With3(3),
+            MockSupertypes2With3(3),
+            MockSupertypes2With3(3),
+        ),
+    )
+
+  def testListSerialization(self):
+    list_original = default_types.List(
+        default_types.Literal(1),
+        default_types.Literal(2),
+        default_types.Literal(3),
+    )
+
+    self.assertEqual(
+        str(list_original), 'List[Literal[1], Literal[2], Literal[3]]'
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(list_original)),
+        list_original,
+    )
+
+  def testTupleSupertype(self):
+    tuple_a = default_types.Tuple(
+        MockSupertypes2With3(1),
+        MockSupertypes2With3(2),
+        MockSupertypes2With3(3),
+    )
+    tuple_b = default_types.Tuple(
+        MockSupertypes2With3(2),
+        MockSupertypes2With3(2),
+        MockSupertypes2With3(2),
+    )
+
+    self.assertEqual(tuple_a, tuple_a.most_specific_common_supertype([]))
+    self.assertIsNone(tuple_a.most_specific_common_supertype([tuple_b]))
+    self.assertEqual(
+        tuple_b.most_specific_common_supertype([tuple_a]),
+        default_types.Tuple(
+            MockSupertypes2With3(3),
+            MockSupertypes2With3(3),
+            MockSupertypes2With3(3),
+        ),
+    )
+
+  def testTupleSerialization(self):
+    tuple_original = default_types.Tuple(
+        default_types.Literal(1),
+        default_types.Literal(2),
+        default_types.Literal(3),
+    )
+
+    self.assertEqual(
+        str(tuple_original), 'Tuple[Literal[1], Literal[2], Literal[3]]'
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(tuple_original)),
+        tuple_original,
+    )
+
+  def testNamedTupleSupertype(self):
+    named_tuple_type = collections.namedtuple('MyNamedTuple', 'x y z')
+    tuple_a = default_types.NamedTuple.from_type_and_attributes(
+        named_tuple_type,
+        (
+            MockSupertypes2With3(1),
+            MockSupertypes2With3(2),
+            MockSupertypes2With3(3),
+        ),
+    )
+    tuple_b = default_types.NamedTuple.from_type_and_attributes(
+        named_tuple_type,
+        (
+            MockSupertypes2With3(2),
+            MockSupertypes2With3(2),
+            MockSupertypes2With3(2),
+        ),
+    )
+
+    self.assertEqual(
+        str(tuple_a),
+        "MyNamedTuple[['x', MockSupertypes2With3], ['y', MockSupertypes2With3],"
+        " ['z', MockSupertypes2With3]]",
+    )
+    self.assertEqual(tuple_a, tuple_a.most_specific_common_supertype([]))
+    self.assertIsNone(tuple_a.most_specific_common_supertype([tuple_b]))
+    self.assertEqual(
+        tuple_b.most_specific_common_supertype([tuple_a]),
+        default_types.NamedTuple.from_type_and_attributes(
+            named_tuple_type,
+            (
+                MockSupertypes2With3(3),
+                MockSupertypes2With3(3),
+                MockSupertypes2With3(3),
+            ),
+        ),
+    )
+
+  def testAttrsSupertype(self):
+    attrs_a = default_types.Attrs.from_type_and_attributes(
+        TestAttrsClass,
+        (
+            MockSupertypes2With3(1),
+            MockSupertypes2With3(2),
+            MockSupertypes2With3(3),
+        ),
+    )
+    attrs_b = default_types.Attrs.from_type_and_attributes(
+        TestAttrsClass,
+        (
+            MockSupertypes2With3(2),
+            MockSupertypes2With3(2),
+            MockSupertypes2With3(2),
+        ),
+    )
+
+    self.assertEqual(
+        str(attrs_a),
+        "TestAttrsClass[['a', MockSupertypes2With3], ['b',"
+        ' MockSupertypes2With3]]',
+    )
+    self.assertEqual(attrs_a, attrs_a.most_specific_common_supertype([]))
+    self.assertIsNone(attrs_a.most_specific_common_supertype([attrs_b]))
+    self.assertEqual(
+        attrs_b.most_specific_common_supertype([attrs_a]),
+        default_types.Attrs.from_type_and_attributes(
+            TestAttrsClass,
+            (
+                MockSupertypes2With3(3),
+                MockSupertypes2With3(3),
+                MockSupertypes2With3(3),
+            ),
+        ),
+    )
+
+  def testDictTypeSubtype(self):
+    dict_type = default_types.Dict
+    dict_a = dict_type(
+        {'a': Mock2AsTopType(1), 'b': Mock2AsTopType(1), 'c': Mock2AsTopType(1)}
+    )
+    dict_b = dict_type(
+        {'a': Mock2AsTopType(2), 'b': Mock2AsTopType(2), 'c': Mock2AsTopType(2)}
+    )
+    dict_c = dict_type({'a': Mock2AsTopType(1), 'b': Mock2AsTopType(1)})
+
+    self.assertTrue(dict_a.is_subtype_of(dict_b))
+    self.assertFalse(dict_c.is_subtype_of(dict_b))
+    self.assertFalse(dict_c.is_subtype_of(dict_a))
+
+  def testDictTypeSupertype(self):
+    dict_type = default_types.Dict
+    dict_a = dict_type({
+        'a': MockSupertypes2With3(1),
+        'b': MockSupertypes2With3(2),
+        'c': MockSupertypes2With3(3),
+    })
+    dict_b = dict_type({
+        'a': MockSupertypes2With3(2),
+        'b': MockSupertypes2With3(2),
+        'c': MockSupertypes2With3(2),
+    })
+
+    self.assertEqual(dict_a, dict_a.most_specific_common_supertype([]))
+    self.assertIsNone(dict_a.most_specific_common_supertype([dict_b]))
+    self.assertEqual(
+        dict_b.most_specific_common_supertype([dict_a]),
+        dict_type({
+            'a': MockSupertypes2With3(3),
+            'b': MockSupertypes2With3(3),
+            'c': MockSupertypes2With3(3),
+        }),
+    )
+
+  def testDictSerialization(self):
+    dict_original = default_types.Dict({
+        'a': default_types.Literal(1),
+        'b': default_types.Literal(2),
+        'c': default_types.Literal(3),
+    })
+
+    self.assertEqual(
+        str(dict_original),
+        "Dict[['a', Literal[1]], ['b', Literal[2]], ['c', Literal[3]]]",
+    )
+    self.assertEqual(
+        serialization.deserialize(serialization.serialize(dict_original)),
+        dict_original,
+    )
+
+  def testListTupleInequality(self):
+    literal = default_types.Literal
+
+    list_a = default_types.List(literal(1), literal(2), literal(3))
+    list_b = default_types.List(literal(1), literal(2), literal(3))
+
+    tuple_a = default_types.Tuple(literal(1), literal(2), literal(3))
+    tuple_b = default_types.Tuple(literal(1), literal(2), literal(3))
+
+    self.assertEqual(list_a, list_b)
+    self.assertEqual(tuple_a, tuple_b)
+    self.assertNotEqual(list_a, tuple_a)
+    self.assertNotEqual(tuple_a, list_a)
+
+  def testDictTypeEquality(self):
+    dict_type = default_types.Dict
+    literal = default_types.Literal
+
+    dict_a = dict_type({literal(1): literal(2), literal(3): literal(4)})
+    dict_b = dict_type({literal(1): literal(2)})
+    dict_c = dict_type({literal(3): literal(4), literal(1): literal(2)})
+
+    self.assertEqual(dict_a, dict_c)
+    self.assertNotEqual(dict_a, dict_b)
+
+  def testCastLazy(self):
+    list_type = default_types.List(
+        default_types.Literal('a'), default_types.Literal('b')
+    )
+    tuple_type = default_types.Tuple(default_types.Literal('c'), list_type)
+    dict_type = default_types.Dict(
+        {'key': tuple_type, 'other_key': list_type}, placeholder_type=dict
+    )
+
+    value = dict_type.placeholder_value(None)
+    casted_value = dict_type.cast(value, None)
+
+    self.assertIs(value, casted_value)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/core/function/trace_type/serialization.proto

@@ -0,0 +1,25 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto2";
+
+package tensorflow.core.function.trace_type.serialization;
+
+import "google/protobuf/any.proto";
+
+// Represents an intermediary representation of a Serializable TraceType.
+message SerializedTraceType {
+  optional google.protobuf.Any representation = 1;
+}

tensorflow/core/function/trace_type/serialization.py

@@ -0,0 +1,100 @@
+# 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.
+# ==============================================================================
+"""Utils for serializing and deserializing TraceTypes."""
+
+import abc
+from typing import Type
+
+from google.protobuf import message
+from tensorflow.core.function.trace_type import serialization_pb2
+
+SerializedTraceType = serialization_pb2.SerializedTraceType
+
+PROTO_CLASS_TO_PY_CLASS = {}
+
+
+class Serializable(metaclass=abc.ABCMeta):
+  """TraceTypes implementing this additional interface are portable."""
+
+  @classmethod
+  @abc.abstractmethod
+  def experimental_type_proto(cls) -> Type[message.Message]:
+    """Returns the unique type of proto associated with this class."""
+    raise NotImplementedError
+
+  @classmethod
+  @abc.abstractmethod
+  def experimental_from_proto(cls, proto: message.Message) -> "Serializable":
+    """Returns an instance based on a proto."""
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def experimental_as_proto(self) -> message.Message:
+    """Returns a proto representing this instance."""
+    raise NotImplementedError
+
+
+def register_serializable(cls: Type[Serializable]):
+  """Registers a Python class to support serialization.
+
+  Only register standard TF types. Custom types should NOT be registered.
+
+  Args:
+    cls: Python class to register.
+  """
+  if cls.experimental_type_proto() in PROTO_CLASS_TO_PY_CLASS:
+    raise ValueError(
+        "Existing Python class " +
+        PROTO_CLASS_TO_PY_CLASS[cls.experimental_type_proto()].__name__ +
+        " already has " + cls.experimental_type_proto().__name__ +
+        " as its associated proto representation. Please ensure " +
+        cls.__name__ + " has a unique proto representation.")
+
+  PROTO_CLASS_TO_PY_CLASS[cls.experimental_type_proto()] = cls
+
+
+def serialize(to_serialize: Serializable) -> SerializedTraceType:
+  """Converts Serializable to a proto SerializedTraceType."""
+
+  if not isinstance(to_serialize, Serializable):
+    raise ValueError("Can not serialize " + type(to_serialize).__name__ +
+                     " since it is not Serializable. For object " +
+                     str(to_serialize))
+  actual_proto = to_serialize.experimental_as_proto()
+
+  if not isinstance(actual_proto, to_serialize.experimental_type_proto()):
+    raise ValueError(
+        type(to_serialize).__name__ +
+        " returned different type of proto than specified by " +
+        "experimental_type_proto()")
+
+  serialized = SerializedTraceType()
+  serialized.representation.Pack(actual_proto)
+  return serialized
+
+
+def deserialize(proto: SerializedTraceType) -> Serializable:
+  """Converts a proto SerializedTraceType to instance of Serializable."""
+  for proto_class in PROTO_CLASS_TO_PY_CLASS:
+    if proto.representation.Is(proto_class.DESCRIPTOR):
+      actual_proto = proto_class()
+      proto.representation.Unpack(actual_proto)
+      return PROTO_CLASS_TO_PY_CLASS[proto_class].experimental_from_proto(
+          actual_proto)
+
+  raise ValueError(
+      "Can not deserialize proto of url: ", proto.representation.type_url,
+      " since no matching Python class could be found. For value ",
+      proto.representation.value)

tensorflow/core/function/trace_type/serialization_test.proto

@@ -0,0 +1,36 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto2";
+
+package tensorflow.core.function.trace_type.serialization_test;
+
+import "tensorflow/core/function/trace_type/serialization.proto";
+
+// Represents a class with just two fields.
+message MyCustomRepresentation {
+  optional int32 index = 1;
+  optional string name = 2;
+}
+
+// Represents a class that is composed of multiple SerializedTraceType objects.
+message MyCompositeRepresentation {
+  repeated tensorflow.core.function.trace_type.serialization.SerializedTraceType
+      elements = 1;
+}
+
+message MyMultiClassRepresentation {
+  optional int32 id = 1;
+}

tensorflow/core/function/trace_type/serialization_test.py

@@ -0,0 +1,239 @@
+# 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 serialization."""
+
+from tensorflow.core.function.trace_type import serialization
+from tensorflow.core.function.trace_type import serialization_test_pb2
+from tensorflow.python.platform import test
+
+
+class MyCustomClass(serialization.Serializable):
+
+  def __init__(self, index, name):
+    self.index = index
+    self.name = name
+
+  @classmethod
+  def experimental_type_proto(cls):
+    return serialization_test_pb2.MyCustomRepresentation
+
+  @classmethod
+  def experimental_from_proto(cls, proto):
+    return MyCustomClass(proto.index, proto.name)
+
+  def experimental_as_proto(self):
+    proto = serialization_test_pb2.MyCustomRepresentation(
+        index=self.index, name=self.name)
+    return proto
+
+
+serialization.register_serializable(MyCustomClass)
+
+
+class MyCompositeClass(serialization.Serializable):
+
+  def __init__(self, *elements):
+    self.elements = elements
+
+  @classmethod
+  def experimental_type_proto(cls):
+    return serialization_test_pb2.MyCompositeRepresentation
+
+  @classmethod
+  def experimental_from_proto(cls, proto):
+    return MyCompositeClass(
+        *[serialization.deserialize(element) for element in proto.elements])
+
+  def experimental_as_proto(self):
+    serialized_elements = [
+        serialization.serialize(element) for element in self.elements
+    ]
+    proto = serialization_test_pb2.MyCompositeRepresentation(
+        elements=serialized_elements)
+    return proto
+
+serialization.register_serializable(MyCompositeClass)
+
+
+class MySerializableSuperClass(serialization.Serializable):
+
+  @classmethod
+  def experimental_type_proto(cls):
+    return serialization_test_pb2.MyMultiClassRepresentation
+
+  @classmethod
+  def experimental_from_proto(cls, proto):
+    if proto.id == 1:
+      return SerializableFromSuperClassOne()
+
+    if proto.id == 2:
+      return SerializableFromSuperClassTwo()
+
+    if proto.id == 3:
+      return SerializableFromSuperClassThree()
+
+    raise NotImplementedError
+
+  def experimental_as_proto(self):
+    if isinstance(self, SerializableFromSuperClassOne):
+      return serialization_test_pb2.MyMultiClassRepresentation(id=1)
+
+    if isinstance(self, SerializableFromSuperClassTwo):
+      return serialization_test_pb2.MyMultiClassRepresentation(id=2)
+
+    if isinstance(self, SerializableFromSuperClassThree):
+      return serialization_test_pb2.MyMultiClassRepresentation(id=3)
+
+    raise NotImplementedError
+
+  def __eq__(self, other):
+    return type(self) is type(other)
+
+serialization.register_serializable(MySerializableSuperClass)
+
+
+class SerializableFromSuperClassOne(MySerializableSuperClass):
+  pass
+
+
+class SerializableFromSuperClassTwo(MySerializableSuperClass):
+  pass
+
+
+class SerializableFromSuperClassThree(MySerializableSuperClass):
+  pass
+
+
+class SerializeTest(test.TestCase):
+
+  def testCustomClassSerialization(self):
+    my_custom = MyCustomClass(1234, "my_name")
+    serialized = serialization.serialize(my_custom)
+
+    self.assertTrue(
+        serialized.representation.Is(
+            serialization_test_pb2.MyCustomRepresentation.DESCRIPTOR))
+
+    proto = serialization_test_pb2.MyCustomRepresentation()
+    serialized.representation.Unpack(proto)
+    self.assertEqual(proto.index, my_custom.index)
+    self.assertEqual(proto.name, my_custom.name)
+
+  def testCustomClassDeserialization(self):
+    original = MyCustomClass(1234, "my_name")
+    serialized = serialization.serialize(original)
+    deserialized = serialization.deserialize(serialized)
+
+    self.assertIsInstance(deserialized, MyCustomClass)
+    self.assertEqual(deserialized.index, original.index)
+    self.assertEqual(deserialized.name, original.name)
+
+  def testCompositeClassSerialization(self):
+    my_composite = MyCompositeClass(
+        MyCustomClass(1, "name_1"), MyCustomClass(2, "name_2"),
+        MyCustomClass(3, "name_3"))
+    serialized = serialization.serialize(my_composite)
+
+    self.assertTrue(
+        serialized.representation.Is(
+            serialization_test_pb2.MyCompositeRepresentation.DESCRIPTOR))
+
+    proto = serialization_test_pb2.MyCompositeRepresentation()
+    serialized.representation.Unpack(proto)
+
+    self.assertEqual(proto.elements[0],
+                     serialization.serialize(MyCustomClass(1, "name_1")))
+    self.assertEqual(proto.elements[1],
+                     serialization.serialize(MyCustomClass(2, "name_2")))
+    self.assertEqual(proto.elements[2],
+                     serialization.serialize(MyCustomClass(3, "name_3")))
+
+  def testCompositeClassDeserialization(self):
+    original = MyCompositeClass(
+        MyCustomClass(1, "name_1"), MyCustomClass(2, "name_2"),
+        MyCustomClass(3, "name_3"))
+    serialized = serialization.serialize(original)
+    deserialized = serialization.deserialize(serialized)
+
+    self.assertIsInstance(deserialized, MyCompositeClass)
+
+    self.assertEqual(deserialized.elements[0].index, 1)
+    self.assertEqual(deserialized.elements[1].index, 2)
+    self.assertEqual(deserialized.elements[2].index, 3)
+
+    self.assertEqual(deserialized.elements[0].name, "name_1")
+    self.assertEqual(deserialized.elements[1].name, "name_2")
+    self.assertEqual(deserialized.elements[2].name, "name_3")
+
+  def testNonUniqueProto(self):
+    class ClassThatReusesProto(serialization.Serializable):
+
+      @classmethod
+      def experimental_type_proto(cls):
+        return serialization_test_pb2.MyCustomRepresentation
+
+      @classmethod
+      def experimental_from_proto(cls, proto):
+        raise NotImplementedError
+
+      def experimental_as_proto(self):
+        raise NotImplementedError
+
+    with self.assertRaisesRegex(
+        ValueError,
+        ("Existing Python class MyCustomClass already has "
+         "MyCustomRepresentation as its associated proto representation. "
+         "Please ensure ClassThatReusesProto has a unique proto representation."
+        )):
+      serialization.register_serializable(ClassThatReusesProto)
+
+  def testWrongProto(self):
+
+    class ClassReturningWrongProto(serialization.Serializable):
+
+      @classmethod
+      def experimental_type_proto(cls):
+        return serialization.SerializedTraceType
+
+      @classmethod
+      def experimental_from_proto(cls, proto):
+        raise NotImplementedError
+
+      def experimental_as_proto(self):
+        return serialization_test_pb2.MyCustomRepresentation()
+
+    with self.assertRaisesRegex(
+        ValueError,
+        ("ClassReturningWrongProto returned different type of proto than "
+         "specified by experimental_type_proto()")):
+      serialization.serialize(ClassReturningWrongProto())
+
+  def testSerializableSuperClass(self):
+    self.assertEqual(
+        serialization.deserialize(
+            serialization.serialize(SerializableFromSuperClassOne())),
+        SerializableFromSuperClassOne())
+    self.assertEqual(
+        serialization.deserialize(
+            serialization.serialize(SerializableFromSuperClassTwo())),
+        SerializableFromSuperClassTwo())
+    self.assertEqual(
+        serialization.deserialize(
+            serialization.serialize(SerializableFromSuperClassThree())),
+        SerializableFromSuperClassThree())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/core/function/trace_type/trace_type_builder.py

@@ -0,0 +1,208 @@
+# 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.
+# ==============================================================================
+"""Utitiles for Cache Key generation based on Function Trace Type."""
+
+import collections.abc
+from typing import Any, Dict, Hashable, Optional
+import weakref
+
+from tensorflow.core.function.trace_type import custom_nest_trace_type
+from tensorflow.core.function.trace_type import default_types
+from tensorflow.core.function.trace_type import util
+from tensorflow.python.types import trace
+from tensorflow.python.util import custom_nest_protocol
+
+
+class InternalTracingContext(trace.TracingContext):
+  """Container for variables and flags shared across TraceType generation."""
+
+  def __init__(self, is_legacy_signature: bool = False):
+    self._global_to_local_id = {}
+    self._alias_id_to_placeholder = {}
+    self._is_legacy_signature = is_legacy_signature
+
+  def alias_global_id(self, global_id: Hashable) -> Hashable:
+    if global_id not in self._global_to_local_id:
+      self._global_to_local_id[global_id] = len(self._global_to_local_id)
+
+    return self._global_to_local_id[global_id]
+
+  def add_placeholder(self, alias_id: Hashable, variable) -> None:
+    self._alias_id_to_placeholder[alias_id] = variable
+
+  def get_placeholder_mapping(self) -> Dict[Hashable, Any]:
+    return self._alias_id_to_placeholder
+
+  @property
+  def is_legacy_signature(self) -> bool:
+    """If the value is from a legacy signature representation.
+
+    Legacy signature representations include tf.function.input_signature and
+    ConcreteFunction.structured_input_signature.
+    """
+    return self._is_legacy_signature
+
+
+class InternalPlaceholderContext(trace.PlaceholderContext):
+  """Container with mappings shared across TraceTypes for placeholder values."""
+
+  def __init__(self,
+               context_graph=None,
+               placeholder_mapping=None,
+               unnest_only=False,
+               with_none_control_dependencies=False,
+               composite_device_name=None):
+    self._alias_id_to_placeholder = placeholder_mapping or {}
+    self._naming_scope = None
+    self._context_graph = context_graph
+    self._unnest_only = unnest_only
+    self._with_none_control_dependencies = with_none_control_dependencies
+    self._composite_device_name = composite_device_name
+
+  def has_placeholder(self, alias_id: Hashable) -> bool:
+    return alias_id in self._alias_id_to_placeholder
+
+  def get_placeholder(self, alias_id: Hashable) -> Hashable:
+    if not self.has_placeholder(alias_id):
+      raise KeyError(f"alias_id: {alias_id} not found in this instance of "
+                     "placeholder context.")
+    return self._alias_id_to_placeholder[alias_id]
+
+  def add_placeholder(self, alias_id: Hashable, placeholder: Hashable) -> None:
+    if alias_id in self._alias_id_to_placeholder:
+      raise KeyError(f"alias id: {alias_id} is already stored in this "
+                     "instance of placeholder context.")
+    self._alias_id_to_placeholder[alias_id] = placeholder
+
+  def update_naming_scope(self, naming_scope: Optional[str]) -> None:
+    self._naming_scope = naming_scope
+
+  @property
+  def naming_scope(self) -> Optional[str]:
+    return self._naming_scope
+
+  @property
+  def context_graph(self):
+    return self._context_graph
+
+  @property
+  def unnest_only(self) -> bool:
+    return self._unnest_only
+
+  @property
+  def with_none_control_dependencies(self) -> bool:
+    return self._with_none_control_dependencies
+
+  @property
+  def composite_device_name(self) -> Any:
+    return self._composite_device_name
+
+
+class InternalCastContext(trace.CastContext):
+  """Default casting behaviors."""
+
+  def __init__(self, allow_specs=False):
+    self._allow_specs = allow_specs
+
+  @property
+  def allow_specs(self) -> bool:
+    """Allow TypeSpecs to be casted (instead of the actual CompositeTensors)."""
+    # Public APIs like get_concrete_function allow users to pass in specs
+    # instead which need to pass through input binding etc.
+    return self._allow_specs
+
+
+def from_value(value: Any,
+               context: trace.TracingContext = None) -> trace.TraceType:
+  """Returns a TraceType corresponding to the value based on the context.
+
+  Args:
+    value: The value to generate a TraceType for.
+    context: The TracingContext to be shared during protocol calls.
+
+  Returns:
+    A TraceType object representing the given value.
+  """
+
+  if context is None:
+    context = InternalTracingContext()
+
+  if context.is_legacy_signature and isinstance(value, trace.TraceType):
+    return value
+  elif isinstance(value, trace.SupportsTracingProtocol):
+    generated_type = value.__tf_tracing_type__(context)
+    if not isinstance(generated_type, trace.TraceType):
+      raise TypeError(
+          "Expected an instance of TraceType for Tracing Protocol call to " +
+          str(value) + " but got " + str(generated_type))
+    return generated_type
+
+  # TODO(b/183107079): Allow these once they're handled properly.
+  if isinstance(value, weakref.ref):
+    raise TypeError(
+        f"weakref input {value} not supported for tf.function."
+    )
+
+  if hasattr(value, "__wrapped__"):
+    return from_value(value.__wrapped__, context)
+
+  if isinstance(value, list):
+    return default_types.List(*(from_value(c, context) for c in value))
+
+  if isinstance(value, tuple):
+    if util.is_namedtuple(value):
+      named_tuple_type = type(value)
+      return default_types.NamedTuple.from_type_and_attributes(
+          named_tuple_type, tuple(from_value(c, context) for c in value))
+    else:
+      return default_types.Tuple(*(from_value(c, context) for c in value))
+
+  if isinstance(value, collections.abc.Mapping):
+    mapping_type = type(value)
+    return default_types.Dict(
+        {k: from_value(value[k], context) for k in value}, mapping_type)
+
+  if util.is_attrs(value):
+    return default_types.Attrs.from_type_and_attributes(
+        type(value),
+        tuple(
+            from_value(getattr(value, a.name), context)
+            for a in value.__attrs_attrs__))
+
+  if util.is_np_ndarray(value):
+    ndarray = value.__array__()
+    return default_types.TENSOR(ndarray.shape, ndarray.dtype)
+
+  if isinstance(value, custom_nest_protocol.CustomNestProtocol):
+    metadata, components = value.__tf_flatten__()
+    return custom_nest_trace_type.CustomNestTraceType(
+        type(value), metadata, tuple(from_value(c, context) for c in components)
+    )
+
+  try:
+    ref = weakref.ref(value)
+    if ref is None:
+      raise TypeError(
+          f"Deleted objects are not valid tf.function arguments, Got {value!r}")
+    else:
+      return default_types.Weakref(ref)
+  except TypeError:
+    try:
+      return default_types.Literal(value)
+    except:
+      raise TypeError(  # pylint: disable=raise-missing-from
+          f"Could not generate a generic TraceType for {value!r}."
+          f"Please verify that it is immutable/hashable. Otherwise, consider "
+          f"implementing the Tracing Protocol for it.")

tensorflow/core/function/trace_type/trace_type_test.py

@@ -0,0 +1,606 @@
+# 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 and benchmarks for the trace_type module."""
+
+import collections
+import dataclasses
+import timeit
+import weakref
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.trace_type import custom_nest_trace_type
+from tensorflow.core.function.trace_type import default_types
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+class TestAttr:
+  """Helps test attrs collections."""
+
+  def __init__(self, name):
+    self.name = name
+
+
+class TestAttrsClass:
+  """Helps test attrs collections."""
+
+  __attrs_attrs__ = (TestAttr('a'), TestAttr('b'))
+
+  def __init__(self, a, b):
+    self.a = a
+    self.b = b
+
+  def __eq__(self, other):
+    return isinstance(
+        other, TestAttrsClass) and self.a == other.a and self.b == other.b
+
+
+class DummyGenericClass:
+  """Helps test memory leaks for GenericType."""
+
+
+@dataclasses.dataclass
+class MaskedTensor:
+  mask: bool
+  value: tensor_lib.Tensor
+
+  def __tf_flatten__(self):
+    metadata = (self.mask,)
+    components = (self.value,)
+    return metadata, components
+
+  @classmethod
+  def __tf_unflatten__(cls, metadata, leaves):
+    mask = metadata[0]
+    value = leaves[0]
+    return MaskedTensor(mask=mask, value=value)
+
+
+class TraceTypeBuilderTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testCompositeAndSpec(self):
+    composite_tensor = ragged_tensor.RaggedTensor.from_row_splits(
+        values=[1, 2, 3], row_splits=[0, 2, 3])
+    spec = ragged_tensor.RaggedTensorSpec([2, None], dtypes.int32)
+
+    self.assertEqual(
+        trace_type.from_value(composite_tensor), trace_type.from_value(spec))
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testVariableAliasing(self):
+    v1 = resource_variable_ops.ResourceVariable([1])
+    v2 = resource_variable_ops.ResourceVariable([1])
+    v3 = resource_variable_ops.ResourceVariable([1])
+    all_unique = trace_type.from_value((v1, v2, v3))
+    all_same = trace_type.from_value((v1, v1, v1))
+    self.assertNotEqual(all_unique, all_same)
+
+    v3 = resource_variable_ops.ResourceVariable([2])
+    v4 = resource_variable_ops.ResourceVariable([2])
+    v5 = resource_variable_ops.ResourceVariable([2])
+    all_unique_again = trace_type.from_value((v3, v4, v5))
+    all_same_again = trace_type.from_value((v4, v4, v4))
+    self.assertEqual(all_unique, all_unique_again)
+    self.assertEqual(all_same, all_same_again)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testTensorEquality(self):
+    context = trace_type.InternalTracingContext()
+    tensor_a = array_ops.zeros([11, 3, 5],
+                               dtype=dtypes.int32).__tf_tracing_type__(context)
+    tensor_b = array_ops.zeros([11, 4, 5],
+                               dtype=dtypes.int32).__tf_tracing_type__(context)
+    tensor_c = array_ops.zeros(
+        [11, 3, 5], dtype=dtypes.float32).__tf_tracing_type__(context)
+    tensor_d = array_ops.ones([11, 3, 5],
+                              dtype=dtypes.int32).__tf_tracing_type__(context)
+
+    self.assertNotEqual(tensor_a, tensor_b)
+    self.assertNotEqual(tensor_a, tensor_c)
+    self.assertNotEqual(tensor_b, tensor_c)
+    self.assertEqual(tensor_a, tensor_d)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testTensorAndSpecEquality(self):
+    context = trace_type.InternalTracingContext()
+    tensor = array_ops.zeros([11, 3, 5],
+                             dtype=dtypes.int32).__tf_tracing_type__(context)
+    spec = tensor_spec.TensorSpec(
+        [11, 3, 5], dtype=dtypes.int32).__tf_tracing_type__(context)
+    spec_with_name = tensor_spec.TensorSpec(
+        [11, 3, 5], dtype=dtypes.int32,
+        name='name').__tf_tracing_type__(context)
+
+    self.assertEqual(tensor, spec)
+    self.assertNotEqual(tensor, spec_with_name)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testTensorShapeUnknown(self):
+    context = trace_type.InternalTracingContext()
+    spec_1 = tensor_spec.TensorSpec(
+        None, dtype=dtypes.int32).__tf_tracing_type__(context)
+    spec_2 = tensor_spec.TensorSpec(
+        None, dtype=dtypes.int32).__tf_tracing_type__(context)
+    self.assertEqual(spec_1, spec_2)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testAttrsTraceTypeGeneration(self):
+    trace_a = trace_type.from_value(TestAttrsClass(1, 2))
+    expected = default_types.Attrs.from_type_and_attributes(
+        TestAttrsClass, (default_types.Literal(1), default_types.Literal(2)))
+    self.assertEqual(trace_a, expected)
+    self.assertTrue(trace_a.is_subtype_of(trace_a))
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testTupleEquality(self):
+    trace_a = trace_type.from_value((1, 2, 3, 4))
+    trace_b = trace_type.from_value((1, 2, 2, 4))
+    trace_c = trace_type.from_value((1, 2, 3))
+    trace_d = trace_type.from_value((1, 2, 3, 4))
+    self.assertNotEqual(trace_a, trace_b)
+    self.assertNotEqual(trace_a, trace_c)
+    self.assertNotEqual(trace_b, trace_c)
+    self.assertEqual(trace_a, trace_d)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testListEquality(self):
+    trace_a = trace_type.from_value([1, 2, 3, 4])
+    trace_b = trace_type.from_value([1, 2, 2, 4])
+    trace_c = trace_type.from_value([1, 2, 3])
+    trace_d = trace_type.from_value([1, 2, 3, 4])
+    self.assertNotEqual(trace_a, trace_b)
+    self.assertNotEqual(trace_a, trace_c)
+    self.assertNotEqual(trace_b, trace_c)
+    self.assertEqual(trace_a, trace_d)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testDictEquality(self):
+    trace_a = trace_type.from_value({1: 2, 3: 4})
+    trace_b = trace_type.from_value({1: 2, 3: 2})
+    trace_c = trace_type.from_value({1: 2, 3: 0})
+    trace_d = trace_type.from_value({3: 4, 1: 2})
+    self.assertNotEqual(trace_a, trace_b)
+    self.assertNotEqual(trace_a, trace_c)
+    self.assertNotEqual(trace_b, trace_c)
+    self.assertEqual(trace_a, trace_d)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testComplexStruct(self):
+    struct = {(1, 2, 3): {(1, 2): {12: 2}}, (3, 2, 3): (2, {2: 3})}
+    trace_a = trace_type.from_value(struct)
+    trace_b = trace_type.from_value(struct)
+    self.assertEqual(trace_a, trace_b)
+    self.assertTrue(trace_a.is_subtype_of(trace_b))
+    self.assertTrue(trace_b.is_subtype_of(trace_a))
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testCustomUnequableTypeSucceeds(self):
+
+    class CustomUnequable:
+
+      def __eq__(self, o):
+        raise ValueError
+
+      def __hash__(self):
+        return 0
+
+    object_a = CustomUnequable()
+    object_b = CustomUnequable()
+    trace_a_1 = trace_type.from_value(object_a)
+    trace_a_2 = trace_type.from_value(object_a)
+    trace_b = trace_type.from_value(object_b)
+    self.assertEqual(trace_a_1, trace_a_2)
+
+    with self.assertRaises(ValueError):
+      trace_a_1.__eq__(trace_b)
+
+    del object_a
+    self.assertNotEqual(trace_a_1, trace_a_2)
+    self.assertNotEqual(trace_a_2, trace_a_1)
+
+    del object_b
+    self.assertNotEqual(trace_a_1, trace_a_2)
+    self.assertNotEqual(trace_a_2, trace_a_1)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testCustomUnhashableTypeFailsGracefully(self):
+
+    class CustomUnhashable:
+
+      def __eq__(self, o):
+        return True
+
+    obj = CustomUnhashable()
+    with self.assertRaisesRegex(
+        TypeError,
+        r'Could not generate a generic TraceType for'):
+      trace_type.from_value(obj)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testGetDefaultPlaceholderValue(self):
+    placeholder_context = trace_type.InternalPlaceholderContext()
+    composite_value = [1, 2, (3, [4, 5]), {6: [7]}, TestAttrsClass(8, (10, 11))]
+    composite_type = trace_type.from_value(composite_value)
+    placeholder_value = composite_type.placeholder_value(placeholder_context)
+    self.assertEqual(composite_value, placeholder_value)
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testWrappedNamedTuple(self):
+    ActualType = collections.namedtuple('ActualType', ['a', 'b', 'c'])
+
+    class MockWrapper(tuple):
+      # Generated through trackable data structures:
+      # //tensorflow/python/trackable/data_structures.py
+      # With design pattern similar to Python functools:
+      # https://docs.python.org/3/library/functools.html?highlight=__wrapped__#functools.update_wrapper
+      __wrapped__ = ActualType(1, 2, 3)
+
+    self.assertEqual(
+        trace_type.from_value(MockWrapper()),
+        trace_type.from_value(ActualType(1, 2, 3)))
+
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
+  def testBadReturnType(self):
+    class MyClass:
+
+      def __tf_tracing_type__(self, _):
+        return 1
+
+    with self.assertRaises(TypeError):
+      trace_type.from_value(MyClass())
+
+
+class CastDefaultTypesTest(test.TestCase, parameterized.TestCase):
+
+  def testLiteral(self):
+    trace_float = default_types.Literal(1.5)
+    ctx = trace_type.InternalCastContext()
+    value = trace_float.cast(1.5, ctx)
+    self.assertEqual(value, 1.5)
+    with self.assertRaises(ValueError):
+      _ = trace_float.cast(1, ctx)
+
+  @parameterized.parameters(list, tuple)
+  def testTupleAndList(self, container_type):
+    foo = (
+        constant_op.constant(1.0, dtypes.float32),
+        constant_op.constant(2.0, dtypes.float32))
+    foo = container_type(foo)
+    trace_foo = trace_type.from_value(foo)
+    bar = (1, 2)
+    bar = container_type(bar)
+    ctx = trace_type.InternalCastContext()
+    value = trace_foo.cast(bar, ctx)
+
+    self.assertIsInstance(value, container_type)
+    self.assertLen(value, len(bar))
+    self.assertSequenceEqual(value, bar)
+    self.assertEqual(value[0].dtype, dtypes.float32)
+    self.assertEqual(value[1].dtype, dtypes.float32)
+
+  @parameterized.parameters(
+      (list, tuple),
+      (tuple, list))
+  def testTupleAndListCannotBeCasted(self, type_a, type_b):
+    foo = (
+        constant_op.constant(1.0, dtypes.float32),
+        constant_op.constant(2.0, dtypes.float32))
+    foo = type_a(foo)
+    trace_foo = trace_type.from_value(foo)
+    bar = (1, 2)
+    bar = type_b(bar)
+    ctx = trace_type.InternalCastContext()
+    with self.assertRaises(AssertionError):
+      _ = trace_foo.cast(bar, ctx)
+
+  def testNamedTuple(self):
+    Foo = collections.namedtuple('Foo', ['x', 'y'])
+    foo = Foo(
+        constant_op.constant(1.0, dtypes.float32),
+        constant_op.constant(2.0, dtypes.float32))
+    trace_foo = trace_type.from_value(foo)
+    bar = Foo(1, 2)
+    ctx = trace_type.InternalCastContext()
+    value = trace_foo.cast(bar, ctx)
+
+    self.assertIsInstance(value, Foo)
+    self.assertLen(value, len(bar))
+    self.assertSequenceEqual(value, bar)
+    self.assertEqual(value[0].dtype, dtypes.float32)
+    self.assertEqual(value[1].dtype, dtypes.float32)
+
+  def testAttrs(self):
+    foo = TestAttrsClass(
+        constant_op.constant(1.0, dtypes.float32),
+        constant_op.constant(2.0, dtypes.float32),)
+    trace_foo = trace_type.from_value(foo)
+    bar = TestAttrsClass(1, 2)
+    ctx = trace_type.InternalCastContext()
+    value = trace_foo.cast(bar, ctx)
+
+    self.assertIsInstance(value, TestAttrsClass)
+    self.assertEqual(value.a.dtype, dtypes.float32)
+    self.assertEqual(value.b.dtype, dtypes.float32)
+
+  def testDict(self):
+    foo = {'x': constant_op.constant(1.0, dtypes.float32),
+           'y': constant_op.constant(2.0, dtypes.float32)}
+    trace_foo = trace_type.from_value(foo)
+    bar = {'x': 1, 'y': 2}
+    ctx = trace_type.InternalCastContext()
+    value = trace_foo.cast(bar, ctx)
+
+    self.assertIsInstance(value, dict)
+    self.assertSequenceEqual(
+        set(value.keys()),
+        set(bar.keys()))
+    self.assertIn('x', value)
+    self.assertIn('y', value)
+    self.assertEqual(value['x'].dtype, dtypes.float32)
+    self.assertEqual(value['y'].dtype, dtypes.float32)
+
+  def testNumpy(self):
+    ndarray = np.array([1, 2, 3])
+    ndarray_type = trace_type.from_value(ndarray)
+    self.assertEqual(
+        ndarray_type, default_types.TENSOR(ndarray.shape, ndarray.dtype)
+    )
+
+  def testWeakrefInput(self):
+    obj = DummyGenericClass()
+    ref = weakref.ref(obj)
+    with self.assertRaisesRegex(TypeError, 'weakref input .* not supported'):
+      trace_type.from_value(ref)
+
+  def testCustomNestCast(self):
+    mt = MaskedTensor(
+        mask=False, value=constant_op.constant([1.0], dtype=dtypes.float32)
+    )
+    mt2 = MaskedTensor(mask=False, value=[2])
+    ctx = trace_type.InternalCastContext()
+
+    trace_mt = trace_type.from_value(mt)
+    self.assertIsInstance(trace_mt, custom_nest_trace_type.CustomNestTraceType)
+
+    mt2_casted = trace_mt.cast(mt2, ctx)
+    self.assertIsInstance(mt2_casted, MaskedTensor)
+    self.assertEqual(mt2_casted.mask, mt2.mask)
+    self.assertEqual(mt2_casted.value.dtype, mt.value.dtype)
+    self.assertAllEqual(mt2_casted.value.shape, mt.value.shape)
+    self.assertAllEqual(mt2_casted.value, mt2.value)
+
+  def testCustomNestFailCastWithWrongMetadata(self):
+    mt = MaskedTensor(mask=False, value=constant_op.constant([1.0]))
+    mt2 = MaskedTensor(mask=True, value=constant_op.constant([1.0]))
+    ctx = trace_type.InternalCastContext()
+    trace_mt = trace_type.from_value(mt)
+    with self.assertRaisesRegex(ValueError, 'Metadata mismatch'):
+      trace_mt.cast(mt2, ctx)
+
+
+class SignatureToTraceTypeTest(test.TestCase):
+
+  def testTensorSpecs(self):
+    self.assertEqual(
+        trace_type.from_value(
+            tensor_spec.TensorSpec(shape=None),
+            trace_type.InternalTracingContext(is_legacy_signature=True)),
+        tensor_spec.TensorSpec(shape=None))
+
+  def testListofTensorSpecs(self):
+    self.assertEqual(
+        trace_type.from_value([
+            tensor_spec.TensorSpec(shape=None),
+            tensor_spec.TensorSpec(shape=None)
+        ], trace_type.InternalTracingContext(is_legacy_signature=True)),
+        default_types.List(
+            tensor_spec.TensorSpec(shape=None),
+            tensor_spec.TensorSpec(shape=None)))
+
+  def testDictofTensorSpecs(self):
+    self.assertEqual(
+        trace_type.from_value(
+            {
+                'a': tensor_spec.TensorSpec(shape=None),
+                'b': tensor_spec.TensorSpec(shape=None)
+            }, trace_type.InternalTracingContext(is_legacy_signature=True)),
+        default_types.Dict({
+            'a': tensor_spec.TensorSpec(shape=None),
+            'b': tensor_spec.TensorSpec(shape=None)
+        }))
+
+
+class TraceTypeMemoryTest(test.TestCase):
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testGeneric(self):
+    trace_type.from_value(1)
+    trace_type.from_value(DummyGenericClass())
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testTensor(self):
+    tensor = array_ops.zeros([10])
+    trace_type.from_value(tensor)
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testTuple(self):
+    trace_type.from_value((1, 2, 3))
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testDict(self):
+    trace_type.from_value({1: 1, 2: 2, 3: 3})
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testList(self):
+    trace_type.from_value([1, 2, 3])
+
+  @test_util.assert_no_new_pyobjects_executing_eagerly()
+  def testAttrs(self):
+    trace_type.from_value(TestAttrsClass(1, 2))
+
+
+class TraceTypeGenerationBenchmark(test.Benchmark):
+
+  def benchmarkTensor(self):
+    shapes = [[1], [2, 19], [5, 11, 24], [4, 5, 9, 23]]
+    tensors = []
+    for s in shapes:
+      tensors.append(array_ops.zeros(s))
+
+    def encode_tensors(tensors):
+      trace_type.from_value(tensors)
+
+    iterations = 100000
+    t = timeit.timeit(lambda: encode_tensors(tensors), number=iterations)
+    self.report_benchmark(
+        name='tensor_cache_key_generation',
+        iters=iterations,
+        wall_time=t,
+        metrics=[{
+            'name': 'tensor_cache_key_generation_avg_ms',
+            'value': t / iterations * 1000
+        }])
+
+  def benchmarkTensorSpec(self):
+    shapes = [[1], [2, 19], [5, 11, 24], [4, 5, 9, 23]]
+    tensor_specs = []
+    for s in shapes:
+      tensor_specs.append(tensor_spec.TensorSpec(s, dtypes.int32))
+
+    def encode_tensor_specs(tensor_specs):
+      trace_type.from_value(tensor_specs)
+
+    iterations = 100000
+    t = timeit.timeit(
+        lambda: encode_tensor_specs(tensor_specs), number=iterations)
+    self.report_benchmark(
+        name='tensor_spec_cache_key_generation',
+        iters=iterations,
+        wall_time=t,
+        metrics=[{
+            'name': 'tensor_spec_cache_key_generation_avg_ms',
+            'value': t / iterations * 1000
+        }])
+
+  def benchmarkVariable(self):
+    var_list = [
+        variables.Variable(1.0),
+        variables.Variable(1),
+        variables.Variable([1])
+    ]
+
+    def encode_variables(var_list):
+      trace_type.from_value(var_list)
+
+    iterations = 10000
+    t = timeit.timeit(lambda: encode_variables(var_list), number=iterations)
+    self.report_benchmark(
+        name='variable_cache_key_generation',
+        iters=iterations,
+        wall_time=t,
+        metrics=[{
+            'name': 'variable_cache_key_generation_avg_ms',
+            'value': t / iterations * 1000
+        }])
+
+  def benchmarkTraceTypeLookup(self):
+
+    @def_function.function
+    def defined(t):
+      return t
+
+    call_arg_list = [
+        1,
+        array_ops.zeros([5, 13]),
+        array_ops.zeros([9, 22, 24]),
+        array_ops.zeros([5, 13, 2])
+    ]
+
+    for c in call_arg_list:
+      defined(c)
+
+    lookup_call_arg = array_ops.zeros([5, 13])
+
+    iterations = 10000
+    t = timeit.timeit(stmt=lambda: defined(lookup_call_arg), number=iterations)
+
+    self.report_benchmark(
+        name='cache_key_lookup',
+        iters=iterations,
+        wall_time=t,
+        metrics=[{
+            'name': 'cache_key_lookup_avg_ms',
+            'value': t / iterations * 1000
+        }])
+
+  def benchmarkNestedStruct(self):
+    struct = {(1, 2, 3): {(1, 2): {12: 2}}, (3, 2, 3): (2, {2: 3})}
+
+    def encode_struct(struct):
+      trace_type.from_value(struct)
+
+    iterations = 100000
+    t = timeit.timeit(lambda: encode_struct(struct), number=iterations)
+    self.report_benchmark(
+        name='nested_struct_cache_key_generation',
+        iters=iterations,
+        wall_time=t,
+        metrics=[{
+            'name': 'nested_struct_cache_key_generation_avg_ms',
+            'value': t / iterations * 1000
+        }])
+
+  def benchmarkFunctionInvocation(self):
+    struct = (variables.Variable(1.0), array_ops.zeros([5, 13]), {
+        'tensor': array_ops.zeros([5, 20]),
+        'variable': variables.Variable(1.0)
+    })
+
+    @def_function.function
+    def defined(t):
+      return t
+
+    defined(struct)  # Get it traced and cached.
+
+    iterations = 10000
+    t = timeit.timeit(lambda: defined(struct), number=iterations)
+    self.report_benchmark(
+        name='function_invocation',
+        iters=iterations,
+        wall_time=t,
+        metrics=[{
+            'name': 'function_invocation_time_avg_ms',
+            'value': t / iterations * 1000
+        }])
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/core/function/trace_type/util.py

@@ -0,0 +1,52 @@
+# 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 the trace_type module."""
+
+from typing import Any, List, Tuple
+
+import numpy as np
+
+
+# TODO(b/225045380): Depend on the abstracted `leaf` lib from 'nest'.
+def is_namedtuple(obj):
+  return hasattr(obj, "_fields") and all(
+      isinstance(field, str) for field in obj._fields)
+
+
+# TODO(b/225045380): Depend on the abstracted `leaf` lib from 'nest'.
+def is_attrs(obj):
+  return hasattr(type(obj), "__attrs_attrs__")
+
+
+# TODO(b/225045380): Depend on the abstracted `leaf` lib from 'nest'.
+def is_np_ndarray(value):
+  return hasattr(value, "__array__") and not (
+      # For legacy reasons we do not automatically promote Numpy strings.
+      isinstance(value, np.str_)
+      # NumPy dtypes have __array__ as unbound methods.
+      or isinstance(value, type))
+
+
+def cast_and_return_whether_casted(
+    trace_types, values, context
+) -> Tuple[List[Any], bool]:
+  did_cast = False
+  casted_values = []
+  for t, v in zip(trace_types, values):
+    casted_v = t.cast(v, context)
+    casted_values.append(casted_v)
+    if casted_v is not v:
+      did_cast = True
+  return casted_values, did_cast

tensorflow/core/function/transform/BUILD

@@ -0,0 +1,62 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/core/function/runtime_client:__subpackages__",
+    ],
+)
+
+licenses(["notice"])
+
+pytype_strict_library(
+    name = "transform",
+    srcs = [
+        "transform.py",
+    ],
+    visibility = ["//smartass/brain/configure:__subpackages__"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/core/framework:function_proto_py",
+        "//tensorflow/core/function/capture:restore_captures",
+        "//tensorflow/core/function/runtime_client:runtime_client_py",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:function_def_to_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:custom_gradient",
+        "//tensorflow/python/ops:default_gradient",
+        "//tensorflow/python/ops:handle_data_util",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/util:compat",
+    ],
+)
+
+py_test(
+    name = "transform_test",
+    srcs = ["transform_test.py"],
+    strict_deps = True,
+    tags = ["no_oss"],  # TODO(b/219089812)
+    deps = [
+        ":transform",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//tensorflow/core/function/testing:test_pass_py",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:custom_gradient",
+        "//tensorflow/python/ops:gradients_impl",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+    ],
+)

tensorflow/core/function/transform/transform.py

@@ -0,0 +1,437 @@
+# 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.
+# ==============================================================================
+"""High level TF Function transformation API."""
+
+from typing import Any, Callable, Iterator, Optional, Union
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.function.capture import restore_captures
+from tensorflow.core.function.runtime_client import runtime_client
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as function_lib
+from tensorflow.python.framework import func_graph as func_graph_module
+from tensorflow.python.framework import function_def_to_graph as function_def_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import custom_gradient as custom_gradient_lib
+from tensorflow.python.ops import default_gradient
+from tensorflow.python.ops import handle_data_util
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.util import compat
+
+_TensorType = Union[ops.EagerTensor, tensor.Tensor]
+_FunctionDefTransformerType = Callable[[function_pb2.FunctionDef], None]
+
+
+def transform_function(
+    f: def_function.Function,
+    inputs: Optional[list[Any]] = None,
+    kw_inputs: Optional[dict[str, Any]] = None,
+    transform_fn: Optional[
+        Union[_FunctionDefTransformerType, list[_FunctionDefTransformerType]]
+    ] = None,
+    mlir_pipeline: Optional[Union[str, list[str]]] = None,
+    nested_fn_transforms: Optional[
+        dict[
+            str,
+            Optional[
+                Union[
+                    _FunctionDefTransformerType,
+                    list[_FunctionDefTransformerType],
+                ]
+            ],
+        ]
+    ] = None,
+    nested_mlir_transforms: Optional[
+        dict[str, Optional[Union[str, list[str]]]]
+    ] = None,
+) -> function_lib.ConcreteFunction:
+  """Applies a transformation to a tf.function to produce a new callable.
+
+  When `transform_fn` is specified, the underlying `FunctionDef` is modified
+  according to the `transform_fn`.
+
+  When `mlir_pipeline` is specified, the underlying `FunctionDef` is converted
+  to an MLIR representation and transformed based on the rules of the
+  `mlir_pipeline`.
+
+  If both are provided, `mlir_pipeline` is applied followed by `transform_fn`.
+
+  Optionally, `transform_fn` could be a list of transformation functions and
+  `mlir_pipeline` could be a list of MLIR transformations. The transformations
+  will be applied in order of the list. For each nested `FunctionDef`, MLIR
+  transformations will be applied before Python function based transformations.
+
+
+  Example:
+  ```python
+  def edit_fn(fndef):
+    for node in fndef.node_def:
+      if node.name == "x_plus_y":
+        node.name = "x_times_y"
+        node.op = "Mul"
+      for idx, inp in enumerate(node.input):
+        if inp == "x_plus_y:z:0":
+          node.input[idx] = "x_times_y:z:0"
+
+  @tf.function(input_signature=[
+      tf.TensorSpec((), dtype=tf.float32),
+      tf.TensorSpec((), dtype=tf.float32)
+  ])
+  def add(x, y):
+    return tf.add(x, y, name="x_plus_y")
+
+  multiply = transform_function(add, transform_fn=edit_fn)
+  assert multiply(1.0, 2.0) == 2.0
+  ```
+
+  Args:
+    f: The target tf.function.
+    inputs: The inputs or input_signature of the tf.function. This does not need
+      to be specified if the `input_signature` was specified in the tf.function
+      decorator.
+    kw_inputs: The keyword inputs of the tf.function. This does not need to be
+      specified if the `input_signature` was specified in the tf.function
+      decorator.
+    transform_fn: A single transformation function or a list of transformation
+      functions to apply on the `FunctionDef`.
+    mlir_pipeline: A single MLIR pass or a list of MLIR passes to transform the
+      `FunctionDef`.
+    nested_fn_transforms: A dict of Python function based transformations to
+      apply on functions in the library of `f`. The keys are the names of the
+      library functions being targeted for transformation.
+    nested_mlir_transforms: A dict of MLIR pass based transformations to apply
+      on functions in the library of `f`. The keys are the names of the library
+      functions being targeted for transformation.
+
+  Returns:
+    The transformed function.
+  """
+
+  # Early exit if no transformations need to be applied.
+  if transform_fn is None and mlir_pipeline is None:
+    return f
+
+  if transform_fn is None:
+    transform_fns = []
+  elif isinstance(transform_fn, list):
+    transform_fns = transform_fn
+  else:
+    transform_fns = [transform_fn]
+
+  if mlir_pipeline is None:
+    mlir_pipelines = []
+  elif isinstance(mlir_pipeline, list):
+    mlir_pipelines = mlir_pipeline
+  else:
+    mlir_pipelines = [mlir_pipeline]
+
+  nested_fn_transforms = (
+      nested_fn_transforms if nested_fn_transforms is not None else {}
+  )
+  nested_mlir_transforms = (
+      nested_mlir_transforms if nested_mlir_transforms is not None else {}
+  )
+
+  # Extract the `ConcreteFunction` from the `tf.function.`
+  if inputs is not None or kw_inputs is not None:
+    inputs = [] if inputs is None else inputs
+    kw_inputs = {} if kw_inputs is None else kw_inputs
+    cf = f.get_concrete_function(*inputs, **kw_inputs)
+  else:
+    cf = f.get_concrete_function()
+
+  # Promote all library functions to the parent scope so that any replicated
+  # functions can also re-use them.
+  graph = ops.get_default_graph()
+  for atomic in cf.graph._functions.values():  # pylint: disable=protected-access
+    graph._add_function_recursive(atomic)  # pylint: disable=protected-access
+
+  # Initialize the `runtime_client`.
+  eager_ctx = runtime_client.GlobalPythonEagerContext()
+  rt = runtime_client.Runtime(eager_ctx)
+
+  # Apply the MLIR passes if provided.
+  for mlir_pipeline in mlir_pipelines:
+    rt.TransformFunction(cf.function_def.signature.name, mlir_pipeline)
+
+  # Get the most up-to-date FunctionDef for the tf.function. This should only
+  # be read after applying any specified mlir_pipelines as they directly
+  # transform the FunctionDef in the runtime.
+  fndef = rt.GetFunctionProto(cf.function_def.signature.name)
+
+  # Apply any transformations if provided.
+  for transform_fn in transform_fns:
+    transform_fn(fndef)
+
+  # Apply a transform to any of the nested AtomicFunctions if
+  # `nested_fn_transforms` or `nested_mlir_transforms` is provided.
+  if nested_fn_transforms or nested_mlir_transforms:
+    nested_functions = cf.graph._functions  # pylint: disable=protected-access
+
+    # Store the new transformed functions.
+    transformed_nested_functions = {}
+
+    # Store a mapping between the old nested function names and the new
+    # transformed function names.
+    nested_transforms_map = {}
+
+    # Transform every nested function specified in `nested_fn_transforms` and
+    # `nested_mlir_transforms`.
+    for atomic_name in (
+        nested_mlir_transforms.keys() | nested_fn_transforms.keys()
+    ):
+      if atomic_name in nested_functions:
+        atomic_transform_fn = nested_fn_transforms.get(atomic_name, [])
+        atomic_mlir_pipeline = nested_mlir_transforms.get(atomic_name, [])
+        transformed_atomic = transform_atomic_function(
+            rt,
+            nested_functions[atomic_name],
+            atomic_transform_fn,
+            atomic_mlir_pipeline,
+        )
+        graph._add_function_recursive(transformed_atomic, overwrite=True)  # pylint: disable=protected-access
+        transformed_atomic_name = compat.as_str(transformed_atomic.name)
+        transformed_nested_functions[transformed_atomic_name] = (
+            transformed_atomic
+        )
+        nested_transforms_map[atomic_name] = transformed_atomic_name
+
+    # Update the `FunctionDef` to map to the newly created EDFs.
+    for node in fndef.node_def:
+      for attr_value in node.attr.values():
+        if attr_value.HasField("func"):
+          attr_value.func.name = nested_transforms_map[attr_value.func.name]
+
+  # Register the updated fndef with the runtime.
+  rt.CreateFunction(fndef)
+
+  # Create a new FuncGraph from the modified FunctionDef.
+  structured_input_signature = cf.structured_input_signature
+  structured_outputs_signature = (
+      func_graph_module.convert_structure_to_signature(cf.structured_outputs)
+  )
+  with graph.as_default():
+    func_graph = function_def_lib.function_def_to_graph(
+        fndef,
+        structured_input_signature=structured_input_signature,
+        structured_outputs=structured_outputs_signature,
+        propagate_device_spec=True,
+    )
+
+  # Set handle data.
+  for i, output in enumerate(cf.outputs):
+    func_graph_output = func_graph.outputs[i]
+    if isinstance(output, tensor.Tensor) and isinstance(
+        func_graph_output, tensor.Tensor
+    ):
+      func_graph_output.set_shape(output.shape)
+      handle_data_util.copy_handle_data(output, func_graph_output)
+
+  # We delete the `_input_shapes` attribute to avoid any intermediate
+  # ShapeInference information from being carried over as the user's
+  # transformations can invalidate them.
+  if "_input_shapes" in fndef.attr:
+    del fndef.attr["_input_shapes"]
+
+  # Replicate custom gradients to the new Graph.
+  with ops.init_scope():
+    _replicate_gradient_functions(cf._func_graph, func_graph)  # pylint: disable=protected-access
+
+  # pylint: disable=protected-access
+  # Get the new ConcreteFunction.
+  updated_cf = function_lib.ConcreteFunction.from_func_graph(
+      func_graph, cf.function_type, attrs=fndef.attr
+  )
+
+  # Set arg_keywords and positional_args
+  updated_cf._arg_keywords = cf._arg_keywords
+  updated_cf._num_positional_args = cf._num_positional_args
+  restore_captures.restore_captures(updated_cf, cf.captured_inputs)
+  # pylint: enable=protected-access
+
+  # Register the ConcreteFunction with the python Graph.
+  if nested_fn_transforms or nested_mlir_transforms:
+    for transformed_atomic in transformed_nested_functions.values():
+      updated_cf.graph._add_function_recursive(  # pylint: disable=protected-access
+          transformed_atomic, overwrite=True
+      )
+  updated_cf.add_to_graph(graph, overwrite=True)
+
+  return updated_cf
+
+
+def transform_atomic_function(
+    rt: runtime_client.Runtime,
+    f: function_lib.AtomicFunction,
+    transform_fn: Union[
+        _FunctionDefTransformerType, list[_FunctionDefTransformerType]
+    ],
+    mlir_pipeline: Union[str, list[str]],
+) -> function_lib.AtomicFunction:
+  """Applies transforms on an AtomicFunction."""
+  transform_fns = (
+      transform_fn if isinstance(transform_fn, list) else [transform_fn]
+  )
+  mlir_pipelines = (
+      mlir_pipeline if isinstance(mlir_pipeline, list) else [mlir_pipeline]
+  )
+  # First apply the MLIR based transformation.
+  for mlir_pipeline in mlir_pipelines:
+    rt.TransformFunction(f.cached_definition.signature.name, mlir_pipeline)
+
+  # Get the `FunctionDef` after MLIR transformation.
+  fndef = rt.GetFunctionProto(f.cached_definition.signature.name)
+
+  # Apply the Python function based transformation.
+  for transform_fn in transform_fns:
+    transform_fn(fndef)
+  rt.CreateFunction(fndef)
+
+  # Generate a new `FuncGraph`
+  graph = ops.get_default_graph()
+  with graph.as_default():
+    func_graph = function_def_lib.function_def_to_graph(
+        fndef,
+        structured_input_signature=f.graph.structured_input_signature,
+        structured_outputs=f.graph.structured_outputs,
+        propagate_device_spec=True,
+    )
+
+  # pylint: disable=protected-access
+  # Ref: third_party/tensorflow/python/ops/control_flow_util_v2.py
+  # Generate a new `AtomicFunction`.
+  atomic = function_lib.from_func_graph(
+      fndef.signature.name,
+      func_graph,
+      fndef.attr,
+      overwrite=True
+  )
+  # pylint: enable=protected-access
+
+  return atomic
+
+
+def _replicate_gradient_functions(
+    original_graph: func_graph_module.FuncGraph,
+    replicated_graph: func_graph_module.FuncGraph,
+) -> None:
+  """Copies over any custom_gradients defined within the original Graph."""
+  seen_ops = set()
+  for gradient_op_type, op in _ops_with_custom_gradients(
+      replicated_graph.get_operations()
+  ):
+    # Soft-cache processed ops so we do not repeat the computation.
+    if gradient_op_type in seen_ops:
+      continue
+    seen_ops.add(gradient_op_type)
+
+    # Lookup the custom_gradient implementation if it exists. Currently all
+    # custom_gradients are stored as python functions in a gradient_registry.
+    # The gradient_registry returns a LookupError when a lookup fails.
+    try:
+      custom_gradient = ops.gradient_registry.lookup(gradient_op_type)
+    except LookupError:
+      continue
+
+    # Convert the custom_gradient to a `ConcreteFunction`. This is done so we
+    # can replicate the custom gradient and update any python captures.
+    try:
+      grad_fn = def_function.function(custom_gradient).get_concrete_function(
+          None, *op.inputs
+      )
+    except Exception:  # pylint: disable=broad-except
+      # TODO(xjun): Figure out why tracing of custom_gradient will fail.
+      tf_logging.exception(
+          f"Error when tracing gradients for {replicated_graph}."
+      )
+      continue
+
+    # Re-bind all captures to values within the replicated graph.
+    remapped_captures = []
+    for capture in grad_fn.captured_inputs:
+      outer_graph, outer_capture = _get_outer_most_capture(
+          original_graph, capture
+      )
+
+      # We only need to re-bind captures originating from the `original_graph`.
+      if outer_graph is not original_graph:
+        continue
+
+      if outer_capture.graph is not outer_graph:
+        raise ValueError(
+            f"Cannot replicate graph: {original_graph}. It utilizes a "
+            f"`tf.custom_gradient` for op: {op} which has a "
+            f"non-replicable capture: {capture}. Consider re-factoring your "
+            "custom_gradient to avoid the capture."
+        )
+
+      remapped_captures.append(
+          replicated_graph.get_tensor_by_name(outer_capture.name)
+      )
+    restore_captures.restore_captures(grad_fn, remapped_captures)
+    new_gradient_op_type = custom_gradient_lib.generate_name()
+    op._set_attr(  # pylint: disable=protected-access
+        "_gradient_op_type",
+        attr_value_pb2.AttrValue(s=compat.as_bytes(new_gradient_op_type)),
+    )
+    ops.RegisterGradient(new_gradient_op_type)(_gen_gradient_func(grad_fn))
+
+
+def _gen_gradient_func(func):
+  """Wraps a ConcreteFunction to be compatible with the gradient registry."""
+
+  def gradient_func(unused_op, *result_grads):
+    result_grads = [
+        x if x is not None else default_gradient.zeros_like(t)
+        for (x, t) in zip(result_grads, func.graph.inputs)
+    ]
+    return func(*result_grads)
+
+  return gradient_func
+
+
+def _get_outer_most_capture(
+    original_graph: func_graph_module.FuncGraph, capture: _TensorType
+) -> tuple[func_graph_module.FuncGraph, _TensorType]:
+  """Tries to find the original captured tensor."""
+  outer_graph = original_graph
+  while outer_graph is not None and not isinstance(capture, ops.EagerTensor):
+    if capture.graph is not outer_graph:
+      outer_graph = outer_graph.outer_graph
+    else:
+      try:
+        capture_index = outer_graph.internal_captures.index(capture)
+      except ValueError:
+        # Capture is a tensor inside the function and is not captured from
+        # another external function
+        break
+      capture = outer_graph.external_captures[capture_index]
+      outer_graph = outer_graph.outer_graph
+
+  return outer_graph, capture
+
+
+def _ops_with_custom_gradients(
+    operations: list[ops.Operation],
+) -> Iterator[tuple[str, ops.Operation]]:
+  """Returns an iterator over ops having custom_gradients."""
+  for op in operations:
+    try:
+      gradient_op_type = op.get_attr("_gradient_op_type")
+    except ValueError:
+      continue
+    yield gradient_op_type, op

tensorflow/core/function/transform/transform_test.py

@@ -0,0 +1,505 @@
+# 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 high-level function transformation API."""
+
+import collections
+
+from absl.testing import parameterized
+
+from tensorflow.core.function.testing import test_pass
+from tensorflow.core.function.transform import transform
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.module import module as module_lib
+from tensorflow.python.ops import custom_gradient
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import load as load_lib
+from tensorflow.python.saved_model import save as save_lib
+
+
+def add_to_multiply(fndef):
+  for node in fndef.node_def:
+    if node.name.endswith("x_plus_y"):
+      node.name = node.name.replace("x_plus_y", "x_times_y")
+      node.op = "Mul"
+    for idx, inp in enumerate(node.input):
+      if inp.endswith("x_plus_y:z:0"):
+        node.input[idx] = inp.replace("x_plus_y", "x_times_y")
+
+
+class Model(module_lib.Module):
+
+  @def_function.function
+  def f(self, x, y, add_2):
+    r = math_ops.add(x, y, name="x_plus_y")
+    if add_2:
+      return r + 2
+    else:
+      return r
+
+
+def apply_transform(f, transform_fn):
+  """Wrapper to apply a transformation on every traced tf.function."""
+
+  @def_function.function
+  def wrapped(*args):
+    updated_cf = transform.transform_function(
+        f, inputs=args, transform_fn=transform_fn)
+    return updated_cf(*args)
+
+  return wrapped
+
+
+class TransformTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="transform",
+          transform_fn=add_to_multiply,
+          mlir_pipeline=None),
+      dict(
+          testcase_name="mlir_pipeline",
+          transform_fn=None,
+          mlir_pipeline="test-pass"),
+      dict(
+          testcase_name="transform_and_mlir_pipeline",
+          transform_fn=add_to_multiply,
+          mlir_pipeline="test-pass"),
+      dict(
+          testcase_name="transform_list",
+          transform_fn=[add_to_multiply],
+          mlir_pipeline=None),
+      dict(
+          testcase_name="mlir_pipeline_list",
+          transform_fn=None,
+          mlir_pipeline=["test-pass"]),
+      dict(
+          testcase_name="transform_list_and_mlir_pipeline_list",
+          transform_fn=[add_to_multiply],
+          mlir_pipeline=["test-pass"]),
+  )
+  @test_util.run_v2_only
+  def test_concrete_function_with(self, transform_fn, mlir_pipeline):
+
+    @def_function.function(input_signature=[
+        tensor_spec.TensorSpec((), dtype=dtypes.float32),
+        tensor_spec.TensorSpec((), dtype=dtypes.float32)
+    ])
+    def f(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    # transfrom f(x, y): x + y -> f(x, y): x * y
+    f = transform.transform_function(
+        f, transform_fn=transform_fn, mlir_pipeline=mlir_pipeline)
+    one = constant_op.constant(1.0)
+    self.assertEqual(f(one, one), 1.0)
+
+    @def_function.function
+    def f2(x, y):
+      z = f(x, y)
+      return math_ops.add(z, 10.0)
+
+    self.assertEqual(f2(one, one), 11.0)
+
+    @def_function.function
+    def f_g(x, y):
+      z = f(x, y)
+      dz_dx, dz_dy = gradients_impl.gradients(z, [x, y])
+      return math_ops.add(dz_dx, dz_dy)
+
+    self.assertEqual(
+        f_g(constant_op.constant(2.0), constant_op.constant(3.0)), 5.0)
+
+  @test_util.run_v2_only
+  def test_transform_with_keywords(self):
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    one = constant_op.constant(1.0)
+
+    # transfrom f(x, y): x + y -> f(x, y): x * y
+    g = transform.transform_function(
+        f,
+        inputs=[one],
+        kw_inputs={"y": one},
+        transform_fn=add_to_multiply,
+        mlir_pipeline="test-pass")
+
+    self.assertEqual(g(one, one), 1.0)
+    self.assertEqual(g(one, y=one), 1.0)
+    self.assertEqual(g(x=one, y=one), 1.0)
+
+  @test_util.run_v2_only
+  def test_transform_with_keywords_only(self):
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    one = constant_op.constant(1.0)
+
+    # transfrom f(x, y): x + y -> f(x, y): x * y
+    g = transform.transform_function(
+        f,
+        inputs=None,
+        kw_inputs={"x": one, "y": one},
+        transform_fn=add_to_multiply,
+        mlir_pipeline="test-pass")
+
+    self.assertEqual(g(one, one), 1.0)
+    self.assertEqual(g(one, y=one), 1.0)
+    self.assertEqual(g(x=one, y=one), 1.0)
+
+  @test_util.run_v2_only
+  def test_function_spec(self):
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.add(x, y, name="x_plus_y")
+
+    args = [1, 1]
+    self.assertEqual(f(*args), 2)
+
+    updated_f = transform.transform_function(
+        f, inputs=args, transform_fn=add_to_multiply)
+    self.assertEqual(updated_f(*args), 1)
+
+    self.assertSequenceAlmostEqual(
+        f.get_concrete_function(
+            *args).pretty_printed_signature().split("\n")[1:],
+        updated_f.pretty_printed_signature().split("\n")[1:])
+
+  @test_util.run_v2_only
+  def test_transform_with_custom_gradients(self):
+
+    @custom_gradient.custom_gradient
+    def add(x, y):
+      e = math_ops.add(x, y, name="x_plus_y")
+
+      # custom gradient that returns gradients of x * y instead of x + y
+      def grad(upstream):
+        dz_dx = y
+        dz_dy = x
+        return upstream * dz_dx, upstream * dz_dy
+
+      return e, grad
+
+    @def_function.function
+    def f(x, y):
+      return add(x, y)
+
+    one = constant_op.constant(1.0)
+    f = transform.transform_function(
+        f, inputs=[one, one], transform_fn=add_to_multiply)
+    self.assertEqual(f(one, one), 1.0)
+
+    @def_function.function
+    def f_g(x, y):
+      z = f(x, y)
+      dz_dx, dz_dy = gradients_impl.gradients(z, [x, y])
+      return math_ops.add(dz_dx, dz_dy)
+
+    self.assertEqual(
+        f_g(constant_op.constant(2.0), constant_op.constant(3.0)), 5.0)
+
+  @test_util.run_v2_only
+  def test_transform_with_nested_function(self):
+
+    @def_function.function
+    def f(x, z):
+
+      @def_function.function
+      def add():
+        i = constant_op.constant(1.0)
+        c = lambda i: math_ops.less(i, 3.0)
+        b = lambda i: (math_ops.add(i, z, name="x_plus_y"))
+        i = while_loop.while_loop_v2(c, b, [i])
+        return i
+
+      y = add()
+      return math_ops.add(x, y, name="x_plus_y")
+
+    one = constant_op.constant(1.0)
+    two = constant_op.constant(2.0)
+    inputs = [one, two]
+
+    # By default only `f` is transformed.
+    updated_f = transform.transform_function(
+        f, inputs=inputs, transform_fn=add_to_multiply)
+    self.assertEqual(updated_f(*inputs), 3.0)  # 1 x (1 + 2) = 3
+
+    # Extract all the functions in `f`'s library that we want to transform.
+    nested_transforms = {}
+    gdef = f.get_concrete_function(*inputs).graph.as_graph_def()
+    for fdef in gdef.library.function:
+      nested_transforms[fdef.signature.name] = add_to_multiply
+
+    # Transform `f` and all of its library functions.
+    updated_f = transform.transform_function(
+        f,
+        inputs=inputs,
+        transform_fn=add_to_multiply,
+        nested_fn_transforms=nested_transforms)
+    self.assertEqual(updated_f(*inputs), 4.0)  # 1 x (1 x 2 x 2) = 4
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="fn",
+          transform_fn=add_to_multiply),
+      dict(
+          testcase_name="mlir",
+          mlir_pipeline="test-pass"),
+      dict(
+          testcase_name="fn_and_mlir",
+          transform_fn=add_to_multiply,
+          mlir_pipeline="test-pass"),
+  )
+  @test_util.run_v2_only
+  def test_nested_python_function_transform_with(self,
+                                                 transform_fn=None,
+                                                 mlir_pipeline=None):
+
+    @def_function.function
+    def f(x, y):
+
+      def inner_add():
+        return math_ops.add(x, y, name="x_plus_y")
+
+      return inner_add()
+
+    inputs = [1.0, 2.0]
+    self.assertEqual(f(*inputs), 3.0)  # 1 + 2 = 3
+
+    # Transform `f`.
+    updated_f = transform.transform_function(
+        f,
+        inputs=inputs,
+        transform_fn=transform_fn,
+        mlir_pipeline=mlir_pipeline)
+    # Nested Python functions should be transformed.
+    self.assertEqual(updated_f(*inputs), 2.0)  # 1 x 2 = 2
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="fn_and_nested_fn",
+          transform_fn=add_to_multiply,
+          nested_fn=True),
+      dict(
+          testcase_name="fn_and_nested_mlir",
+          transform_fn=add_to_multiply,
+          nested_mlir=True),
+      dict(
+          testcase_name="mlir_and_nested_fn",
+          mlir_pipeline="test-pass",
+          nested_fn=True),
+      dict(
+          testcase_name="mlir_and_nested_mlir",
+          mlir_pipeline="test-pass",
+          nested_mlir=True),
+      dict(
+          testcase_name="mlir_fn_and_nested_mlir_fn",
+          transform_fn=add_to_multiply,
+          mlir_pipeline="test-pass",
+          nested_fn=True,
+          nested_mlir=True),
+  )
+  @test_util.run_v2_only
+  def test_nested_transform_with(self,
+                                 transform_fn=None,
+                                 mlir_pipeline=None,
+                                 nested_fn=False,
+                                 nested_mlir=False):
+
+    @def_function.function
+    def f(x, y, z):
+
+      @def_function.function
+      def inner_add():
+        return math_ops.add(y, z, name="x_plus_y")
+
+      return math_ops.add(x, inner_add(), name="x_plus_y")
+
+    # 1, 2, 4 are picked so the following combinations create different results.
+    # 1 + (2 + 4) = 7
+    # 1 + (2 * 4) = 9
+    # 1 * (2 + 4) = 6
+    # 1 * (2 * 4) = 8
+    inputs = [1.0, 2.0, 4.0]
+    self.assertEqual(f(*inputs), 7.0)  # 1 + (2 + 4) = 7
+
+    # Extract all the functions in `f`'s library that we want to transform.
+    nested_fn_transforms = {}
+    nested_mlir_transforms = {}
+
+    cf = f.get_concrete_function(*inputs)
+    gdef = cf.graph.as_graph_def()
+    for fdef in gdef.library.function:
+      fdef_name = fdef.signature.name
+      if nested_fn:
+        nested_fn_transforms[fdef_name] = add_to_multiply
+      if nested_mlir:
+        nested_mlir_transforms[fdef_name] = "test-pass"
+
+    # Transform `f` and all of its library functions.
+    updated_f = transform.transform_function(
+        f,
+        inputs=inputs,
+        transform_fn=transform_fn,
+        mlir_pipeline=mlir_pipeline,
+        nested_fn_transforms=nested_fn_transforms,
+        nested_mlir_transforms=nested_mlir_transforms)
+    self.assertEqual(updated_f(*inputs), 8.0)  # 1 x (2 x 4) = 8
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="fn_and_nested_fn",
+          transform_fn=add_to_multiply,
+          nested_fn=True),
+      dict(
+          testcase_name="fn_and_nested_mlir",
+          transform_fn=add_to_multiply,
+          nested_mlir=True),
+      dict(
+          testcase_name="mlir_and_nested_fn",
+          mlir_pipeline="test-pass",
+          nested_fn=True),
+      dict(
+          testcase_name="mlir_and_nested_mlir",
+          mlir_pipeline="test-pass",
+          nested_mlir=True),
+      dict(
+          testcase_name="mlir_fn_and_nested_mlir_fn",
+          transform_fn=add_to_multiply,
+          mlir_pipeline="test-pass",
+          nested_fn=True,
+          nested_mlir=True),
+  )
+  @test_util.run_v2_only
+  def test_nested_transform_in_tf_function_with(self,
+                                                transform_fn=None,
+                                                mlir_pipeline=None,
+                                                nested_fn=False,
+                                                nested_mlir=False):
+
+    @def_function.function
+    def g(x, y, z):
+
+      @def_function.function
+      def f(x, y, z):
+
+        @def_function.function
+        def inner_add():
+          return math_ops.add(y, z, name="x_plus_y")
+
+        return math_ops.add(x, inner_add(), name="x_plus_y")
+
+      nested_fn_transforms = {}
+      nested_mlir_transforms = {}
+
+      cf = f.get_concrete_function(*inputs)
+      gdef = cf.graph.as_graph_def()
+      for fdef in gdef.library.function:
+        fdef_name = fdef.signature.name
+        if nested_fn:
+          nested_fn_transforms[fdef_name] = add_to_multiply
+        if nested_mlir:
+          nested_mlir_transforms[fdef_name] = "test-pass"
+
+      updated_f = transform.transform_function(
+          f,
+          inputs=inputs,
+          transform_fn=transform_fn,
+          mlir_pipeline=mlir_pipeline,
+          nested_fn_transforms=nested_fn_transforms,
+          nested_mlir_transforms=nested_mlir_transforms)
+
+      return updated_f(x, y, z)
+
+    inputs = [1.0, 2.0, 4.0]
+    graph_def = g.get_concrete_function(*inputs).graph.as_graph_def()
+    # Confirm all "AddV2" nodes in the library functions of graph_def are
+    # transformed to "Mul".
+    ops = collections.Counter()
+    for fdef in graph_def.library.function:
+      for node in fdef.node_def:
+        ops[node.op] += 1
+
+    self.assertNotIn("AddV2", ops)
+    self.assertEqual(ops["Mul"], 2)
+
+  @test_util.run_v2_only
+  def test_save_transform_for_all_signatures(self):
+    m = Model()
+    # Originally f does addition
+    self.assertEqual(
+        m.f(
+            constant_op.constant(1, dtypes.int32),
+            constant_op.constant(2, dtypes.int32),
+            constant_op.constant(False, dtypes.bool)), 3)
+
+    self.assertEqual(
+        m.f(
+            constant_op.constant(1, dtypes.int32),
+            constant_op.constant(2, dtypes.int32),
+            constant_op.constant(True, dtypes.bool)), 5)
+
+    # Transform every input signature of f to a multiply
+    m.f = apply_transform(m.f, add_to_multiply)
+
+    # Validate arbitrary signatures.
+    self.assertEqual(
+        m.f(
+            constant_op.constant(1, dtypes.int32),
+            constant_op.constant(2, dtypes.int32),
+            constant_op.constant(False, dtypes.bool)), 2)
+    self.assertEqual(
+        m.f(
+            constant_op.constant(1, dtypes.int32),
+            constant_op.constant(2, dtypes.int32),
+            constant_op.constant(True, dtypes.bool)), 4)
+    self.assertEqual(
+        m.f(
+            constant_op.constant(1.0, dtypes.float32),
+            constant_op.constant(2.0, dtypes.float32),
+            constant_op.constant(False, dtypes.bool)), (2.0))
+
+    # Save and restore the model.
+    save_lib.save(m, "/tmp/testing_model")
+    m_loaded = load_lib.load("/tmp/testing_model")
+
+    # Validate the restored model.
+    self.assertEqual(
+        m_loaded.f(
+            constant_op.constant(1, dtypes.int32),
+            constant_op.constant(2, dtypes.int32),
+            constant_op.constant(False, dtypes.bool)), 2)
+    self.assertEqual(
+        m_loaded.f(
+            constant_op.constant(1.1, dtypes.float32),
+            constant_op.constant(2.0, dtypes.float32),
+            constant_op.constant(True, dtypes.bool)), (4.2))
+
+
+if __name__ == "__main__":
+  test_pass.RegisterTestPass()
+  test.main()