chore: import upstream snapshot with attribution

tensorflow/python/framework/_dtypes.pyi

@@ -0,0 +1,41 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class DType:
+    def __init__(self, arg0: object) -> None: ...
+    def __hash__(self) -> int: ...
+    def __int__(self) -> int: ...
+    @property
+    def as_datatype_enum(self) -> int: ...
+    @property
+    def is_bool(self) -> bool: ...
+    @property
+    def is_complex(self) -> bool: ...
+    @property
+    def is_floating(self) -> bool: ...
+    @property
+    def is_integer(self) -> bool: ...
+    @property
+    def is_numeric(self) -> bool: ...
+    @property
+    def is_numpy_compatible(self) -> bool: ...
+    @property
+    def is_quantized(self) -> bool: ...
+    @property
+    def is_unsigned(self) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def size(self) -> int: ...

tensorflow/python/framework/_errors_test_helper.pyi

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

tensorflow/python/framework/_op_def_library_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 process_inputs(arg0: str, arg1: int, arg2: dict) -> tuple: ...

tensorflow/python/framework/_op_def_registry.pyi

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

tensorflow/python/framework/_op_def_util.pyi

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

tensorflow/python/framework/_proto_comparators.pyi

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

tensorflow/python/framework/_py_context_manager.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 test_py_context_manager(arg0: object, arg1: object) -> object: ...

tensorflow/python/framework/_python_memory_checker_helper.pyi

@@ -0,0 +1,18 @@
+# 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 Callable
+
+def mark_stack_trace_and_call(arg0: Callable) -> None: ...

tensorflow/python/framework/_pywrap_python_api_dispatcher.pyi

@@ -0,0 +1,58 @@
+# 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
+
+MATCH: MatchType
+MATCH_DISPATCHABLE: MatchType
+NO_MATCH: MatchType
+
+class MatchType:
+    __members__: ClassVar[dict] = ...  # read-only
+    MATCH: ClassVar[MatchType] = ...
+    MATCH_DISPATCHABLE: ClassVar[MatchType] = ...
+    NO_MATCH: ClassVar[MatchType] = ...
+    __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: ...
+
+class PySignatureChecker:
+    def __init__(self, arg0: list[tuple[int, PyTypeChecker]]) -> None: ...
+    def CheckCanonicalizedArgs(self, arg0: tuple) -> bool: ...
+
+class PyTypeChecker:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def Check(self, arg0: object) -> MatchType: ...
+    def cache_size(self) -> int: ...
+    def cost(self) -> int: ...
+
+class PythonAPIDispatcher:
+    def __init__(self, arg0: str, arg1: list[str], arg2: object) -> None: ...
+    def Dispatch(self, arg0: object, arg1: object) -> object: ...
+    def Register(self, arg0: PySignatureChecker, arg1: object) -> None: ...
+    def Unregister(self, arg0: object) -> None: ...
+
+def MakeInstanceChecker(*args) -> PyTypeChecker: ...
+def MakeListChecker(arg0: PyTypeChecker) -> PyTypeChecker: ...
+def MakeUnionChecker(arg0: list[PyTypeChecker]) -> PyTypeChecker: ...
+def register_dispatchable_type(arg0: object) -> object: ...

tensorflow/python/framework/_pywrap_python_api_info.pyi

@@ -0,0 +1,32 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class InferredAttributes:
+    def __init__(self, *args, **kwargs) -> None: ...
+    @property
+    def lengths(self) -> list[int]: ...
+    @property
+    def type_lists(self): ...
+    @property
+    def types(self): ...
+
+class PythonAPIInfo:
+    def __init__(self, arg0: str) -> None: ...
+    def DebugInfo(self) -> str: ...
+    def InferredLengthAttrs(self) -> list[str]: ...
+    def InferredTypeAttrs(self) -> list[str]: ...
+    def InferredTypeListAttrs(self) -> list[str]: ...
+    def InitializeFromParamSpecs(self, arg0: dict[str, str], arg1: dict[str, str], arg2: list[str], arg3: object) -> None: ...
+    def InitializeFromRegisteredOp(self, arg0: str) -> None: ...

tensorflow/python/framework/_pywrap_python_api_parameter_converter.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 Convert(*args, **kwargs): ...

tensorflow/python/framework/_pywrap_python_op_gen.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 GetPythonWrappers(arg0: bytes) -> bytes: ...

tensorflow/python/framework/_pywrap_python_tensor_converter.pyi

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

tensorflow/python/framework/_test_metrics_util.pyi

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

tensorflow/python/framework/auto_control_deps.py

@@ -0,0 +1,651 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""AutomaticControlDependencies and related functionality."""
+
+import collections
+import enum
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import auto_control_deps_utils as utils
+from tensorflow.python.framework import dtypes as dtypes_module
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import registry
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.util import nest
+from tensorflow.python.util import object_identity
+from tensorflow.python.util import tf_decorator
+
+# LINT.IfChange
+# Op types that should not run in program order, e.g. because they need to run
+# asynchronously to avoid deadlock.
+
+ASYNC_STATEFUL_OPS = frozenset((
+    "CollectiveGather",
+    "CollectiveReduce",
+    "CollectiveBcastSend",
+    "CollectiveBcastSendV2",
+    "CollectiveBcastRecv",
+    "CollectiveBcastRecvV2",
+    "NcclAllReduce",
+    # We do not add "Send" here since we want it to be added as a control output
+    # in order to avoid being pruned.
+    "Recv",
+    "CollectiveInitializeCommunicator",
+    "CollectiveAssignGroupV2",
+))
+
+LEGACY_RANDOM_OPS = frozenset((
+    # These may be used in variable initializers -- thus their execution should
+    # not be dependent on other stateful operations.  This is because although
+    # according to program order, tf.Variables may be created in sequence,
+    # their initialization happens outside of the program order (specifically,
+    # in graph mode their initialization happens by calling a grouped
+    # initializer operation or in eager mode, where initialization is lifted
+    # out of the tf.function and executed the first time the function is
+    # executed).
+    #
+    # Unless there is a specific dependency between the initializers
+    # themselves (e.g. one initializer depends on a Variable whose value depends
+    # on another initializer), the initialization can happen in any order so
+    # long as it's before the associated Variable read operations.
+    #
+    # Note that in general the randomness of legacy random operations is only
+    # guaranteed by providing a graph-level and op-level seed (and ordering of
+    # the same op across multiple iterations of a while_loop is specifically not
+    # guaranteed; see the discussion below).
+    #
+    # There is a possible race condition inside while_loop where the same
+    # random OpKernel instantiation is reused across multiple steps
+    # of the loop.  Since legacy Random OpKernels have an internal rng state,
+    # automatic dependency tracking across loop steps would likely
+    # fix this race; and for that case this denylist is problematic.
+    # However, since automatic dependency tracking inside while loops is not
+    # currently supported, and there are no other examples of OpKernel reuse
+    # (each OpKernel is associated with a unique op in graph mode),
+    # this denylist has no effect on the aforementioned behavior.
+    #
+    # TODO(ebrevdo,skyewm): Modify the check against this denylist to
+    # only occur when the op is inside a "variable initialization scope"; and
+    # add proper autodeps inside while_loops that respects this updated check.
+    "RandomUniform",
+    "RandomUniformInt",
+    "RandomStandardNormal",
+    "ParameterizedTruncatedNormal",
+    "TruncatedNormal",
+    "RandomShuffle",
+    "Multinomial",
+    "RandomGamma",
+    "RandomGammaGrad",
+    "RandomPoisson",
+    "RandomPoissonV2",
+))
+
+MUST_RUN_ORDER_INSENSITIVE_STATEFUL_OPS = frozenset((
+    "InfeedEnqueue",
+    "InfeedEnqueueTuple",
+    "EnqueueTPUEmbeddingSparseBatch",
+    "EnqueueTPUEmbeddingIntegerBatch",
+    "EnqueueTPUEmbeddingSparseTensorBatch",
+    "EnqueueTPUEmbeddingRaggedTensorBatch",
+    "EnqueueTPUEmbeddingArbitraryTensorBatch",
+    "DynamicEnqueueTPUEmbeddingArbitraryTensorBatch",
+))
+
+# These ops are order-insensitive ans should in theory run, but at the moment
+# they either always have the necessary data dependencies, or have workarounds
+# in existing code that would break when adding new control deps. This
+# inconsistency should be eventually fixed, but it would be more effective to
+# retire the list instead.
+SKIPPED_ORDER_INSENSITIVE_STATEFUL_OPS = frozenset((
+    "CudnnRNN",
+    "CudnnRNNBackprop",
+    "CudnnRNNV2",
+    "CudnnRNNV3",
+    "CudnnRNNBackpropV2",
+    "CudnnRNNBackpropV3",
+    "RestoreV2",
+    "SaveV2",
+))
+# LINT.ThenChange(//tensorflow/core/grappler/optimizers/function_optimizer.cc)
+
+# Op types that are marked as stateless, but should be allowlisted to add auto
+# control dependencies.
+_ALLOWLIST_STATELESS_OPS = [
+    # As TPU collective ops are blocking, if there are more than one collective
+    # op in the function, we need to make sure different collectives ops are
+    # scheduled in certain orders. Otherwise if at the same time all the
+    # replicas are launching different collective ops/programs, it may cause
+    # deadlock.
+    "AllToAll",
+    "CrossReplicaSum",
+    "CollectivePermute",
+]
+
+
+def op_is_stateful(op):
+  # pylint: disable=protected-access
+  ret = ((op._is_stateful and
+          ((op.type not in ASYNC_STATEFUL_OPS) and
+           (op.type not in LEGACY_RANDOM_OPS) and
+           (op.type not in SKIPPED_ORDER_INSENSITIVE_STATEFUL_OPS))) or
+         (op.type in _ALLOWLIST_STATELESS_OPS))
+  return ret
+
+
+class ResourceType(enum.Enum):
+  READ_ONLY = "read-only"
+  READ_WRITE = "read-write"
+
+
+def collective_manager_ids_from_op(op):
+  """Returns CollectiveManager ID from the op if one exists, else None.
+
+  CollectiveManager adds collective and no_op operations tagged with an ID,
+  unique to the manager object. This function extracts that ID, or None, if the
+  node was not generated by a CollectiveManager.
+
+  Args:
+    op: `Operation` to get the collective manager ID from.
+
+  Returns:
+    List of CollectiveManager IDs used by the op.
+  """
+  if op.type == "CollectiveReduce":
+    try:
+      return [op.get_attr("_collective_manager_id")]
+    except ValueError:
+      pass
+  elif op.type == "StatefulPartitionedCall":
+    try:
+      return op.get_attr(utils.COLLECTIVE_MANAGER_IDS)
+    except ValueError:
+      pass
+  return []
+
+
+class AutomaticControlDependencies(object):
+  """Context manager to automatically add control dependencies.
+
+  Code under this context manager will act as if a sensible set of control
+  dependencies were present. More specifically:
+    1. All stateful ops in the scope will execute (with the exception of ops in
+       ASYNC_STATEFUL_OPS and LEGACY_RANDOM_OPS)
+    2. Stateful ops which modify the same resource will execute in program order
+
+  Note: creating variables in an automatic control dependencies context is not
+  supported (the value of the variables will never change as they will keep
+  getting reinitialized).
+
+  NOT THREAD SAFE
+  """
+
+  def __init__(self):
+    self._returned_tensors = object_identity.ObjectIdentitySet()
+    self.ops_which_must_run = set()
+    self._independent_ops = []
+
+  def mark_as_return(self, tensor):
+    """Acts like identity but marks the `Tensor` as a return value.
+
+    This will possibly return a copy of the `Tensor`. Usage:
+
+    ```
+      with AutomaticControlDependencies() as a:
+       ...
+       t = a.mark_as_return(t)
+      _ = ...(t...)  # i.e. it's safe to use t here
+    ```
+
+    Args:
+      tensor: the `Tensor` to be marked
+
+    Returns:
+      a copy of the `Tensor`.
+    """
+    if isinstance(tensor, indexed_slices.IndexedSlices):
+      values = array_ops.identity(tensor.values)
+      indices = array_ops.identity(tensor.indices)
+      self._returned_tensors.add(indices)
+      self._returned_tensors.add(values)
+      return indexed_slices.IndexedSlices(
+          values, indices, dense_shape=tensor.dense_shape)
+    elif isinstance(tensor, sparse_tensor.SparseTensor):
+      values = array_ops.identity(tensor.values)
+      indices = array_ops.identity(tensor.indices)
+      self._returned_tensors.add(indices)
+      self._returned_tensors.add(values)
+      return sparse_tensor.SparseTensor(
+          indices, values, dense_shape=tensor.dense_shape)
+    elif isinstance(tensor, tensor_array_ops.TensorArray):
+      flow = array_ops.identity(tensor.flow)
+      self._returned_tensors.add(flow)
+      return tensor_array_ops.build_ta_with_new_flow(tensor, flow)
+    # We want to make the return values depend on the stateful operations, but
+    # we don't want to introduce a cycle, so we make the return value the result
+    # of a new identity operation that the stateful operations definitely don't
+    # depend on.
+    tensor = array_ops.identity(tensor)
+    self._returned_tensors.add(tensor)
+    return tensor
+
+  def run_independently(self, op):
+    """Marks the given op as independent.
+
+    Overrides any other rule for the op.
+
+    Independent ops are guaranteed to execute before the return values, but
+    are allowed to run in parallel with everything else. Use in programs which
+    can guarantee that an op has side effects that don't affect any other op.
+
+    Args:
+      op: An operation
+    """
+    self._independent_ops.append(op)
+    op._set_attr("_independent_side_effects", attr_value_pb2.AttrValue(b=True))  # pylint: disable=protected-access
+
+  def __enter__(self):
+    if context.executing_eagerly():
+      return self
+    # This code assumes no other thread is adding ops to the graph while
+    # we're adding ops to the graph.
+    # TODO(apassos): Fix this by locking the graph or using a temporary
+    # graph (but that would mess up devices and collections at least,
+    # probably other things as well).
+    g = ops.get_default_graph()
+    self._graph = g
+    g._add_control_dependencies = True  # pylint: disable=protected-access
+    g.experimental_acd_manager = self
+    self._n_operations = g.num_operations()
+    return self
+
+  def _process_switch(self, switch_op, ops_which_must_run,
+                      last_write_to_resource, merge_for_resource):
+    """Processes a switch node for a resource input.
+
+    When tensorflow creates a cond, it creates a control flow context for each
+    branch of the cond. Each external tensor accessed by that branch is routed
+    through a switch op, which gets created in the graph _after_ the op which
+    uses that tensor get created.
+
+    If the resource comes from another switch op we process that one first.
+
+    _process_switch creates a corresponding merge node for the switch node. This
+    merge node is added to the outer control flow context of the switch
+    node. We also ensure that:
+
+      1. The switch node executes after the previous op which used the resource
+         tensor
+
+      2. Any op which uses a resource output of the switch node executes before
+         the merge for the switch node.
+
+      3. The next op which uses the input resource to the switch node (which
+         might be another switch node for the other branch of the conditional)
+         will execute after the merge node is done.
+
+      4. The merge node is marked as must_run so it will run even if no
+         subsequent operation uses the resource.
+
+    Args:
+      switch_op: the switch op to be processed
+      ops_which_must_run: the set of ops which must run
+      last_write_to_resource: map from resource tensor to last op updating
+        it
+      merge_for_resource: map from resource tensor to merge which must follow
+        all usages of it.
+    """
+    # pylint: disable=protected-access
+    inp = switch_op.inputs[0]
+    input_id = ops.tensor_id(inp)
+    if inp.dtype == dtypes_module.resource and inp.op.type == "Switch":
+      self._process_switch(inp.op, ops_which_must_run, last_write_to_resource,
+                           merge_for_resource)
+    output = switch_op.outputs[0]
+    output_id = ops.tensor_id(output)
+    if output_id in merge_for_resource:
+      return
+    new_merge = control_flow_ops.merge(
+        switch_op.outputs, name="artificial_merge")
+    new_merge[0].op._control_flow_context = (
+        switch_op._control_flow_context.outer_context)
+    # Ensures the merge always runs
+    ops_which_must_run.add(new_merge[0].op)
+    if input_id in last_write_to_resource:
+      # Ensures the switch executes after the previous op using the resource.
+      switch_op._add_control_input(last_write_to_resource[input_id])
+    # Ensure the next op outside the cond happens after the merge.
+    last_write_to_resource[input_id] = new_merge[0].op
+    if input_id in merge_for_resource:
+      merge_for_resource[input_id]._add_control_input(new_merge[0].op)
+    for o in switch_op.outputs:
+      # Ensures the merge will execute after all ops inside the cond
+      merge_for_resource[ops.tensor_id(o)] = new_merge[0].op
+
+  def __exit__(self, unused_type, unused_value, unused_traceback):
+    # pylint: disable=protected-access
+    if context.executing_eagerly():
+      return
+
+    if self._graph is not ops.get_default_graph():
+      raise RuntimeError(
+          "Within the automatic control dependency context, the default graph"
+          f" cannot change. Upon entry it was {self._graph}, but on exit it"
+          f" changed to {ops.get_default_graph()}")
+
+    outer_graph = getattr(self._graph, "outer_graph", None)
+    if outer_graph is not None:
+      self._graph._add_control_dependencies = outer_graph._add_control_dependencies
+    else:
+      self._graph._add_control_dependencies = False
+    self._graph.experimental_acd_manager = None
+
+    # map from resource tensor to the last op which wrote to it
+    last_write_to_resource = {}
+    # map from resource tensor to the list of reads from it since the last
+    # write or since the beginning of the function.
+    reads_since_last_write_to_resource = collections.defaultdict(list)
+    # CollectiveManager manager_ids within a particular function call should not
+    # be needed outside of that function call. So we keep them separate (though
+    # the general idea of the maps is the same, in the future, we'll need to
+    # correctly thread the control output outside).
+    # Map from collective manager scope to the last op which used it
+    collective_manager_scopes_opened = {}
+    collective_manager_scopes_used = {}
+    # set of conditional and loop exits
+    ops_which_must_run = set()
+    # merge which must depend on ops which use this resource
+    merge_for_resource = {}
+
+    new_operations = self._graph.get_operations()[self._n_operations:]
+
+    # Ensures that uses of resource tensors get serialized properly and all
+    # execute. This is done by keeping a map from resource tensor to the last op
+    # in graph-construction order which used it (last_write_to_resource).
+    #
+    # Conditionals are written in TensorFlow such that every external tensor
+    # accessed in the conditional goes through a switch op and every return
+    # tensor (it's guaranteed that there will be at least one) goes through a
+    # merge op.
+    #
+    # To handle conditionals, switches are handled in a special way (see
+    # comments for _process_switch). Merge nodes created by TF's conditional
+    # logic (as opposed to by _process_switch) are forced to run and also get a
+    # control dependency added to them to ensure all stateful ops inside their
+    # control flow context run.
+    #
+    # We also ensure that if an op is using a resource output by a switch node
+    # (that is, a resource tensor for which there's a value in
+    # merge_for_resource) this op will run before the merge for that resource.
+    #
+    # We try to add control inputs to nodes respecting their control flow
+    # contexts to avoid dead nodes propagating everywhere and leading to
+    # "retval[0] doesn't have value" errors. If a node gets a control dependency
+    # on a dead node (i.e. a note from an untaken control flow branch) that node
+    # will be marked as dead unless it's a merge node.
+    #
+    # TODO(apassos): serialize non-resource-taking stateful ops as well, and
+    # test that it works. Support while loops. Support init_scope escaping from
+    # this.
+    for op in new_operations:
+      # TODO(apassos) make this code safely support while loops.
+      if control_flow_util.IsInWhileLoop(op):
+        continue
+      control_inputs = set()
+
+      if op.type in MUST_RUN_ORDER_INSENSITIVE_STATEFUL_OPS:
+        # This will add it to self._independent_ops, but also mark it with an
+        # attribute.
+        self.run_independently(op)
+
+      if op in self._independent_ops:
+        ops_which_must_run.add(op)
+        continue
+
+      # Ensure stateful ops run.
+      # Read-only ops are added to control outputs if the read value is
+      # consumed. This covers the case when the read value is returned from
+      # the function since that goes through a tf.identity in mark_as_return.
+      if ((op_def_registry.get(op.type) is None) or
+          (op_is_stateful(op) and
+           (op.type not in utils.RESOURCE_READ_OPS or
+            any(output.consumers() for output in op.outputs)))):
+        ops_which_must_run.add(op)
+
+      # Make a note of all opened manager_ids.
+      if op.type == "NoOp":
+        try:
+          collective_manager_scopes_opened[op.get_attr(
+              "_collective_manager_id")] = op
+        except ValueError:
+          pass
+      # Ignore switches (they're handled separately)
+      if op.type == "Switch" and op.inputs[0].dtype == dtypes_module.resource:
+        continue
+      # Make merges trigger all other computation which must run
+      # TODO(mdan): Don't do this. Write a transform to chains instead.
+      # See core/common_runtime/control_flow_deps_to_chains.cc.
+      if op.type == "Merge":
+        for o in ops_which_must_run:
+          op._add_control_input(o)
+          for inp in o.inputs:
+            input_id = ops.tensor_id(inp)
+            if input_id in last_write_to_resource:
+              last_write_to_resource[input_id] = op
+        ops_which_must_run = set([op])
+        continue
+
+      resource_inputs = set()
+      # Check for any resource inputs. If we find any, we update control_inputs
+      # and last_write_to_resource.
+      for inp, resource_type in _get_resource_inputs(op):
+        is_read = resource_type == ResourceType.READ_ONLY
+        input_id = ops.tensor_id(inp)
+
+        # If the op receives the same resource tensor twice as an input, we skip
+        # to avoid the op getting a control dependency on itself.
+        if input_id in resource_inputs:
+          continue
+
+        resource_inputs.add(input_id)
+        # Deal with switches, finally.
+        if inp.op.type == "Switch":
+          self._process_switch(inp.op, ops_which_must_run,
+                               last_write_to_resource, merge_for_resource)
+        is_building_function = op.graph.building_function
+        # Ensure uses of resources are serialized
+        if input_id in last_write_to_resource:
+          if is_building_function or (
+              last_write_to_resource[input_id]._control_flow_context
+              is op._control_flow_context):
+            control_inputs.add(last_write_to_resource[input_id])
+        # Ensure merges happen after the closing of a cond block
+        if input_id in merge_for_resource:
+          merge_for_resource[input_id]._add_control_input(op)
+        if is_read:
+          reads_since_last_write_to_resource[input_id].append(op)
+        else:
+          control_inputs.update(reads_since_last_write_to_resource[input_id])
+          reads_since_last_write_to_resource[input_id] = []
+          last_write_to_resource[input_id] = op
+
+      if (op_is_stateful(op) and not resource_inputs
+          and op._control_flow_context is None):
+        if None in last_write_to_resource:
+          op._add_control_input(last_write_to_resource[None])
+        last_write_to_resource[None] = op
+
+      # Ensure ordering of collective ops
+      manager_ids = collective_manager_ids_from_op(op)
+      for manager_id in manager_ids:
+        if manager_id in collective_manager_scopes_opened:
+          # Chain this function call if the scope was opened.
+          op._add_control_input(collective_manager_scopes_opened[manager_id])
+          collective_manager_scopes_opened[manager_id] = op
+        else:
+          # If this op is in a scope not created here, create a chain starting
+          # at this op.
+          if manager_id in collective_manager_scopes_used:
+            op._add_control_input(collective_manager_scopes_used[manager_id])
+          collective_manager_scopes_used[manager_id] = op
+
+      if control_inputs and not is_building_function:
+        control_inputs = [
+            c for c in control_inputs
+            if c._control_flow_context is op._control_flow_context
+        ]
+
+      op._add_control_inputs(control_inputs)
+
+    # Ensure all ops which must run do run
+    self.ops_which_must_run.update(ops_which_must_run)
+
+    control_output_op = None
+    for idx, r in enumerate(
+        nest.flatten(list(self._returned_tensors), expand_composites=True)):
+      if self.ops_which_must_run:
+        updated_ops_which_must_run = []
+        if r.graph.building_function:
+          # There may be many stateful ops in the graph. Adding them as
+          # control inputs to each function output could create excessive
+          # control edges in the graph. Thus we create an intermediate No-op to
+          # chain the control dependencies between stateful ops and function
+          # outputs.
+          if idx == 0:
+            control_output_op = control_flow_ops.no_op()
+            control_output_op._add_control_inputs(self.ops_which_must_run)
+          updated_ops_which_must_run = [control_output_op]
+        else:
+          updated_ops_which_must_run = [
+              o for o in self.ops_which_must_run
+              if o._control_flow_context is r.op._control_flow_context
+          ]
+        r.op._add_control_inputs(updated_ops_which_must_run)
+
+    self.collective_manager_ids_used = collective_manager_scopes_used
+
+
+_acd_resource_resolvers_registry = registry.Registry("acd_resource_resolvers")
+
+
+def register_acd_resource_resolver(f):
+  """Register a function for resolving resources touched by an op.
+
+  `f` is called for every Operation added in the ACD context with the op's
+  original resource reads and writes. `f` is expected to update the sets of
+  resource reads and writes in-place and return True if it updated either of the
+  sets, False otherwise.
+
+  Example:
+  @register_acd_resource_resolver
+  def identity_resolver(op, resource_reads, resource_writes):
+    # op: The `Operation` being processed by ACD currently.
+    # resource_reads: An `ObjectIdentitySet` of read-only resources.
+    # resource_writes: An `ObjectIdentitySet` of read-write resources.
+    def update(resource_inputs):
+      to_remove = []
+      to_add = []
+      for resource in resource_inputs:
+        if resource.op.type == "Identity":
+          to_remove.append(resource)
+          to_add.extend(resource.op.inputs)
+      for t in to_remove:
+        resource_inputs.discard(t)
+      resource_inputs.update(to_add)
+      return to_add or to_remove
+    return update(resource_reads) or update(resource_writes)
+
+  Args:
+    f: Python function with signature
+    (Operation, ObjectIdentitySet, ObjectIdentitySet) -> bool
+
+  Returns:
+    The function `f` after adding it to the registry.
+  """
+  _acd_resource_resolvers_registry.register(f)
+  return f
+
+
+@register_acd_resource_resolver
+def _identity_resolver(op, resource_reads, resource_writes):
+  """Replaces Identity output with its input in resource_inputs."""
+  del op
+  def update(resource_inputs):
+    to_remove = []
+    to_add = []
+    for resource in resource_inputs:
+      if resource.op.type == "Identity":
+        to_remove.append(resource)
+        to_add.extend(resource.op.inputs)
+    for t in to_remove:
+      resource_inputs.discard(t)
+    resource_inputs.update(to_add)
+    return to_add or to_remove
+
+  return update(resource_reads) or update(resource_writes)
+
+
+def _get_resource_inputs(op):
+  """Returns an iterable of resources touched by this `op`."""
+  reads, writes = utils.get_read_write_resource_inputs(op)
+  saturated = False
+  while not saturated:
+    saturated = True
+    for key in _acd_resource_resolvers_registry.list():
+      # Resolvers should return true if they are updating the list of
+      # resource_inputs.
+      # TODO(srbs): An alternate would be to just compare the old and new set
+      # but that may not be as fast.
+      updated = _acd_resource_resolvers_registry.lookup(key)(op, reads, writes)
+      if updated:
+        # Conservatively remove any resources from `reads` that are also writes.
+        reads = reads.difference(writes)
+      saturated = saturated and not updated
+
+  # Note: A resource handle that is not written to is treated as read-only. We
+  # don't have a special way of denoting an unused resource.
+  for t in reads:
+    yield (t, ResourceType.READ_ONLY)
+  for t in writes:
+    yield (t, ResourceType.READ_WRITE)
+
+
+def automatic_control_dependencies(f):
+  """Wraps f to automatically insert control dependencies.
+
+  The inserted dependencies ensure that:
+    1. All stateful ops in f run when the result of f runs
+    2. Updates to the same resources happen in order.
+
+  Args:
+    f: the function to be wrapped.
+
+  Returns:
+    The wrapped function.
+  """
+
+  def wrapper(*args, **kwargs):
+    with AutomaticControlDependencies() as a:
+      result = f(*args, **kwargs)
+      result_flat = [a.mark_as_return(t) for t in nest.flatten(result)]
+      return nest.pack_sequence_as(result, result_flat)
+
+  return tf_decorator.make_decorator(f, wrapper)

tensorflow/python/framework/auto_control_deps_test.py

@@ -0,0 +1,997 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import itertools
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import auto_control_deps as acd
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_switch_case
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import gen_sendrecv_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+from tensorflow.python.training import adam
+from tensorflow.python.training import momentum
+
+
+class AutomaticControlDependenciesTest(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.must_run_order_insensitive_stateful_ops = (
+        acd.MUST_RUN_ORDER_INSENSITIVE_STATEFUL_OPS)
+
+  def tearDown(self):
+    acd.MUST_RUN_ORDER_INSENSITIVE_STATEFUL_OPS = (
+        self.must_run_order_insensitive_stateful_ops)
+    super().tearDown()
+
+  def testBasic(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies() as c:
+        v.assign(v + 1)
+        v.assign(2 * v)
+        val = v.read_value()
+        val = c.mark_as_return(val)
+      self.assertAllEqual(val, 4.0)
+
+  def testUnorderedOpsRunInParallel(self):
+    acd.MUST_RUN_ORDER_INSENSITIVE_STATEFUL_OPS |= frozenset(("EagerPyFunc",))
+
+    side_effects = []
+
+    def side_effect_one(x):
+      side_effects.append(1)
+      return x
+
+    def side_effect_two(x):
+      side_effects.append(2)
+      return x
+
+    @def_function.function
+    def f():
+      script_ops.eager_py_func(side_effect_one, [1], [dtypes.int32])
+      script_ops.eager_py_func(side_effect_two, [1], [dtypes.int32])
+      return 1
+
+    side_effects = []
+    self.evaluate(f())
+
+    self.assertSetEqual(set(side_effects), set((1, 2)))
+
+  def testIndependentOpsRunInParallel(self):
+    v = resource_variable_ops.ResourceVariable(1)
+    self.evaluate(variables.global_variables_initializer())
+
+    @def_function.function
+    def f():
+      gen_resource_variable_ops.assign_variable_op(v.handle, 1)
+      ops.get_default_graph().experimental_acd_manager.run_independently(
+          gen_resource_variable_ops.assign_variable_op(v.handle, 2))
+
+    # A function with two identical ops, should cause a data race in most
+    # conditions.
+    var_values = set()
+    for _ in range(10000):
+      self.evaluate(f())
+      var_values.add(
+          self.evaluate(
+              resource_variable_ops.read_variable_op(v.handle, dtypes.int32)))
+    # With regular control dependencies, the function should always run the
+    # first assign first, and the value 1 should never be seen.
+    # With run_independently, assign 1 and 2 are run in parallel. Thus, when f
+    # is run large number of times, we see both 1 and 2 values assigned to
+    # variable v.
+    self.assertSetEqual(var_values, set((1, 2)))
+
+  def testIndependentOpsInLoop(self):
+    v = resource_variable_ops.ResourceVariable(0)
+    self.evaluate(variables.global_variables_initializer())
+
+    @def_function.function
+    def f():
+      for i in math_ops.range(3):
+        ops.get_default_graph().experimental_acd_manager.run_independently(
+            gen_resource_variable_ops.assign_variable_op(v.handle, i))
+
+    self.evaluate(f())
+    # TODO(mdan): Find a more robust way to test in loops.
+    self.assertEqual(
+        self.evaluate(
+            resource_variable_ops.read_variable_op(v.handle, dtypes.int32)), 2)
+
+  def testNoControlDepsBetweenVariableReads(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies():
+        read_op1 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        read_op2 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+      self.assertNotIn(read_op1, read_op2.control_inputs)
+      self.assertNotIn(read_op2, read_op1.control_inputs)
+
+  def testVariableReadThenWrite(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies():
+        read_op1 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        read_op2 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        assign_op = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+      # Writes should have control deps from "all" reads since last write
+      # or start of the code block.
+      self.assertIn(read_op1, assign_op.control_inputs)
+      self.assertIn(read_op2, assign_op.control_inputs)
+      # There should be no control deps between reads.
+      self.assertNotIn(read_op1, read_op2.control_inputs)
+      self.assertNotIn(read_op2, read_op1.control_inputs)
+
+  def testVariableWriteThenRead(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies():
+        assign_op = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+        read_op1 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        read_op2 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+      # Reads should have a control dep from the last write.
+      self.assertIn(assign_op, read_op1.control_inputs)
+      self.assertIn(assign_op, read_op2.control_inputs)
+      # There should be no control deps between reads.
+      self.assertNotIn(read_op1, read_op2.control_inputs)
+      self.assertNotIn(read_op2, read_op1.control_inputs)
+
+  def testIdentityPassThrough(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies():
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        identity_handle = gen_array_ops.identity(v.handle)
+        assign_op2 = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+        read_op = gen_resource_variable_ops.read_variable_op(
+            identity_handle, v.dtype).op
+      # Read should have a control dep from second last write even
+      # with Identity applied to resource.
+      self.assertIn(assign_op2, read_op.control_inputs)
+
+  def testVariableReadsInOpsWithMustRun(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies() as c:
+        read_op = gen_resource_variable_ops.read_variable_op(v.handle,
+                                                             v.dtype).op
+        # Read ops get added to control outputs only if they have consumers.
+        c.mark_as_return(read_op.outputs[0])
+      self.assertIn(read_op, c.ops_which_must_run)
+
+  def testVariableMultipleReadsAndWrites(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies() as c:
+        # 2 reads -> 2 writes -> 2 reads -> 2 writes.
+        read_op1 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        read_op2 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        assign_op1 = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+        assign_op2 = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+        read_op3 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        read_op4 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        assign_op3 = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+        assign_op4 = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+        # Read ops get added to control outputs only if they have consumers.
+        c.mark_as_return(read_op1.outputs[0])
+        c.mark_as_return(read_op2.outputs[0])
+        c.mark_as_return(read_op3.outputs[0])
+        c.mark_as_return(read_op4.outputs[0])
+
+      # Verify the control edges.
+      self.assertIn(read_op1, assign_op1.control_inputs)
+      self.assertIn(read_op2, assign_op1.control_inputs)
+      self.assertIn(assign_op1, assign_op2.control_inputs)
+      self.assertIn(assign_op2, read_op3.control_inputs)
+      self.assertIn(assign_op2, read_op4.control_inputs)
+      self.assertIn(read_op3, assign_op3.control_inputs)
+      self.assertIn(read_op4, assign_op3.control_inputs)
+      self.assertIn(assign_op3, assign_op4.control_inputs)
+
+      # There should be no control deps between reads.
+      read_ops = [read_op1, read_op2, read_op3, read_op4]
+      for src_op, tgt_op in itertools.product(read_ops, read_ops):
+        self.assertNotIn(src_op, tgt_op.control_inputs)
+
+      # Reads must be in `ops_which_must_run`.
+      self.assertIn(read_op1, c.ops_which_must_run)
+      self.assertIn(read_op2, c.ops_which_must_run)
+      self.assertIn(read_op3, c.ops_which_must_run)
+      self.assertIn(read_op4, c.ops_which_must_run)
+      # Last write must be in `ops_which_must_run`.
+      self.assertIn(assign_op4, c.ops_which_must_run)
+
+  def testSendInOpsWithMustRun(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies() as c:
+        send_op = gen_sendrecv_ops.send(v, "x", "/", 0, "/")
+
+      # Send must be in `ops_which_must_run`.
+      self.assertIn(send_op, c.ops_which_must_run)
+
+  def _testVariableReadInFunctionalOp(self, build_functional_op, op_type):
+    v = resource_variable_ops.ResourceVariable(1.0)
+    self.evaluate(variables.global_variables_initializer())
+
+    @def_function.function
+    def read_var_in_while():
+      gen_resource_variable_ops.read_variable_op(
+          v.handle, v.dtype, name="read1")
+
+      result = build_functional_op(v)
+      gen_resource_variable_ops.read_variable_op(
+          v.handle, v.dtype, name="read2")
+      gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+      return result
+
+    func_graph = read_var_in_while.get_concrete_function().graph
+    assert len(func_graph.inputs) == 1
+
+    def get_op(op_type, sub_name):
+      operations = [
+          op for op in func_graph.get_operations()
+          if op.type == op_type and sub_name in op.name
+      ]
+      assert len(operations) == 1
+      return operations[0]
+
+    read1 = get_op("ReadVariableOp", "read1")
+    functional_op = get_op(op_type, "")
+    read2 = get_op("ReadVariableOp", "read2")
+    assign_op = get_op("AssignVariableOp", "")
+    # Since the functional op only has reads, previous reads e.g. read1 do not\
+    # have a control edge to it and next future reads e.g. read2 do not have a
+    # control edge from it.
+    self.assertNotIn(read1, functional_op.control_inputs)
+    self.assertNotIn(functional_op, read2.control_inputs)
+    self.assertIn(read1, assign_op.control_inputs)
+    self.assertIn(read2, assign_op.control_inputs)
+    self.assertIn(functional_op, assign_op.control_inputs)
+
+  def testVariableReadInWhileLoop(self):
+
+    def build_functional_op(v):
+
+      def body(_):
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return while_loop.while_loop(
+          lambda i: True, body, [0.0], maximum_iterations=1)
+
+    self._testVariableReadInFunctionalOp(build_functional_op, "While")
+
+  def testVariableReadInCondTrueBranch(self):
+
+    def build_functional_op(v):
+
+      def then_branch():
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      def else_branch():
+        return array_ops.zeros([], v.dtype)
+
+      return cond.cond(
+          constant_op.constant(True), then_branch, else_branch)
+
+    self._testVariableReadInFunctionalOp(build_functional_op, "If")
+
+  def testVariableReadInCondFalseBranch(self):
+
+    def build_functional_op(v):
+
+      def then_branch():
+        return array_ops.zeros([], v.dtype)
+
+      def else_branch():
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return cond.cond(
+          constant_op.constant(False), then_branch, else_branch)
+
+    self._testVariableReadInFunctionalOp(build_functional_op, "If")
+
+  def testVariableReadInCaseBranch0(self):
+
+    def build_functional_op(v):
+
+      def branch0():
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      def branch1():
+        return array_ops.zeros([], v.dtype)
+
+      return control_flow_switch_case.switch_case(
+          constant_op.constant(0), [branch0, branch1])
+
+    self._testVariableReadInFunctionalOp(build_functional_op, "Case")
+
+  def testVariableReadInCaseBranch1(self):
+
+    def build_functional_op(v):
+
+      def branch0():
+        return array_ops.zeros([], v.dtype)
+
+      def branch1():
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return control_flow_switch_case.switch_case(
+          constant_op.constant(0), [branch0, branch1])
+
+    self._testVariableReadInFunctionalOp(build_functional_op, "Case")
+
+  def testVariableReadInFunction(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_read():
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return fn_with_read()
+
+    self._testVariableReadInFunctionalOp(build_functional_op,
+                                         "StatefulPartitionedCall")
+
+  def testVariableReadInNestedFunction(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_read():
+
+        @def_function.function
+        def inner_fn():
+          return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+        return inner_fn()
+
+      return fn_with_read()
+
+    self._testVariableReadInFunctionalOp(build_functional_op,
+                                         "StatefulPartitionedCall")
+
+  def testVariableReadInWhileInInnerFunc(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_read():
+
+        @def_function.function
+        def inner_fn():
+
+          def body(_):
+            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+          return while_loop.while_loop(
+              lambda i: True, body, [0.0], maximum_iterations=1)
+
+        return inner_fn()
+
+      return fn_with_read()
+
+    self._testVariableReadInFunctionalOp(build_functional_op,
+                                         "StatefulPartitionedCall")
+
+  def testVariableReadInCondInInnerFunc(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_read():
+
+        @def_function.function
+        def inner_fn():
+
+          def then_branch():
+            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+          def else_branch():
+            return array_ops.zeros([], v.dtype)
+
+          return cond.cond(
+              constant_op.constant(True), then_branch, else_branch)
+
+        return inner_fn()
+
+      return fn_with_read()
+
+    self._testVariableReadInFunctionalOp(build_functional_op,
+                                         "StatefulPartitionedCall")
+
+  def _testVariableWriteInFunctionalOp(self, build_functional_op, op_type):
+    v = resource_variable_ops.ResourceVariable(1.0)
+    self.evaluate(variables.global_variables_initializer())
+
+    @def_function.function
+    def write_var_in_while():
+      gen_resource_variable_ops.read_variable_op(
+          v.handle, v.dtype, name="read1")
+
+      result = build_functional_op(v)
+      gen_resource_variable_ops.read_variable_op(
+          v.handle, v.dtype, name="read2")
+      gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+      return result
+
+    func_graph = write_var_in_while.get_concrete_function().graph
+    assert len(func_graph.inputs) == 1
+
+    def get_op(op_type, sub_name):
+      operations = [
+          op for op in func_graph.get_operations()
+          if op.type == op_type and sub_name in op.name
+      ]
+      assert len(operations) == 1
+      return operations[0]
+
+    read1 = get_op("ReadVariableOp", "read1")
+    functional_op = get_op(op_type, "")
+    read2 = get_op("ReadVariableOp", "read2")
+    assign_op = get_op("AssignVariableOp", "")
+    # Since the While has writes, it has control edges from previous reads
+    # e.g. `read1` and to future reads(`read2`) and writes(`assign_op`).
+    self.assertIn(read1, functional_op.control_inputs)
+    self.assertIn(functional_op, read2.control_inputs)
+    self.assertIn(read2, assign_op.control_inputs)
+    self.assertIn(functional_op, assign_op.control_inputs)
+
+  def testVariableWriteInWhileLoop(self):
+
+    def build_functional_op(v):
+
+      def body(_):
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return while_loop.while_loop(
+          lambda i: True, body, [0.0], maximum_iterations=1)
+
+    self._testVariableWriteInFunctionalOp(build_functional_op, "While")
+
+  def testVariableWriteInCondTrueBranch(self):
+
+    def build_functional_op(v):
+
+      def then_branch():
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      def else_branch():
+        return array_ops.zeros([], v.dtype)
+
+      return cond.cond(
+          constant_op.constant(True), then_branch, else_branch)
+
+    self._testVariableWriteInFunctionalOp(build_functional_op, "If")
+
+  def testVariableWriteInCondFalseBranch(self):
+
+    def build_functional_op(v):
+
+      def then_branch():
+        return array_ops.zeros([], v.dtype)
+
+      def else_branch():
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return cond.cond(
+          constant_op.constant(False), then_branch, else_branch)
+
+    self._testVariableWriteInFunctionalOp(build_functional_op, "If")
+
+  def testVariableWriteInCaseBranch0(self):
+
+    def build_functional_op(v):
+
+      def branch0():
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      def branch1():
+        return array_ops.zeros([], v.dtype)
+
+      return control_flow_switch_case.switch_case(
+          constant_op.constant(0), [branch0, branch1])
+
+    self._testVariableWriteInFunctionalOp(build_functional_op, "Case")
+
+  def testVariableWriteInCaseBranch1(self):
+
+    def build_functional_op(v):
+
+      def branch0():
+        return array_ops.zeros([], v.dtype)
+
+      def branch1():
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return control_flow_switch_case.switch_case(
+          constant_op.constant(0), [branch0, branch1])
+
+    self._testVariableWriteInFunctionalOp(build_functional_op, "Case")
+
+  def testVariableWriteInFunction(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_write():
+        gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+        return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+      return fn_with_write()
+
+    self._testVariableWriteInFunctionalOp(build_functional_op,
+                                          "StatefulPartitionedCall")
+
+  def testVariableWriteInNestedFunction(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_write():
+
+        @def_function.function
+        def inner_fn():
+          gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+          return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+        return inner_fn()
+
+      return fn_with_write()
+
+    self._testVariableWriteInFunctionalOp(build_functional_op,
+                                          "StatefulPartitionedCall")
+
+  def testVariableWriteInWhileInInnerFunc(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_write():
+
+        @def_function.function
+        def inner_fn():
+
+          def body(_):
+            gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+          return while_loop.while_loop(
+              lambda i: True, body, [0.0], maximum_iterations=1)
+
+        return inner_fn()
+
+      return fn_with_write()
+
+    self._testVariableWriteInFunctionalOp(build_functional_op,
+                                          "StatefulPartitionedCall")
+
+  def testVariableWriteInCondInInnerFunc(self):
+
+    def build_functional_op(v):
+
+      @def_function.function
+      def fn_with_write():
+
+        @def_function.function
+        def inner_fn():
+
+          def then_branch():
+            gen_resource_variable_ops.assign_variable_op(v.handle, v + 1)
+            return gen_resource_variable_ops.read_variable_op(v.handle, v.dtype)
+
+          def else_branch():
+            return array_ops.zeros([], v.dtype)
+
+          return cond.cond(
+              constant_op.constant(True), then_branch, else_branch)
+
+        return inner_fn()
+
+      return fn_with_write()
+
+    self._testVariableWriteInFunctionalOp(build_functional_op,
+                                          "StatefulPartitionedCall")
+
+  @test_util.run_v1_only("b/120545219")
+  def testCondMustRun(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      p = array_ops.placeholder(dtype=dtypes.bool)
+      with acd.AutomaticControlDependencies() as c:
+
+        def true_fn():
+          v.assign(v + 1)
+          return 0.0
+
+        def false_fn():
+          v.assign(v + 4)
+          return 1.0
+
+        cond.cond(p, true_fn, false_fn)
+        val = v.read_value()
+        val = c.mark_as_return(val)
+      self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0)
+      self.assertAllEqual(val.eval(feed_dict={p: True}), 6.0)
+
+  @test_util.run_v1_only("b/120545219")
+  def testCondMustRunSeparateRead(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      p = array_ops.placeholder(dtype=dtypes.bool)
+      with acd.AutomaticControlDependencies() as c:
+
+        def true_fn():
+          v.assign(v + 1)
+          return 0.0
+
+        def false_fn():
+          v.assign(v + 4)
+          return 1.0
+
+        cond.cond(p, true_fn, false_fn)
+        one = constant_op.constant(1.0)
+        one = c.mark_as_return(one)
+      one.eval(feed_dict={p: False})
+      self.assertAllEqual(v.read_value(), 5.0)
+      one.eval(feed_dict={p: True})
+      self.assertAllEqual(v.read_value(), 6.0)
+
+  @test_util.run_v1_only("b/120545219")
+  def testCondNested(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      p = array_ops.placeholder(dtype=dtypes.bool)
+      q = array_ops.placeholder(dtype=dtypes.bool)
+      with acd.AutomaticControlDependencies() as c:
+
+        def true_fn():
+          v.assign(v + 1, name="true")
+          return 1.0
+
+        def false_fn():
+
+          def inner_true_fn():
+            v.assign(v * 2, name="false_true")
+            return 2.0
+
+          def inner_false_fn():
+            v.assign(v * 3, name="false_false")
+            return 3.0
+
+          cond.cond(q, inner_true_fn, inner_false_fn)
+          return 1.0
+
+        cond.cond(p, true_fn, false_fn)
+        with ops.name_scope("final"):
+          val = v.read_value()
+        val = c.mark_as_return(val)
+      self.assertAllEqual(val.eval(feed_dict={p: False, q: False}), 3.0)
+      self.assertAllEqual(val.eval(feed_dict={p: False, q: True}), 6.0)
+      self.assertAllEqual(val.eval(feed_dict={p: True, q: True}), 7.0)
+      self.assertAllEqual(val.eval(feed_dict={p: True, q: False}), 8.0)
+
+  @test_util.run_v1_only("b/120545219")
+  def testCondOneBranch(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      p = array_ops.placeholder(dtype=dtypes.bool)
+      with acd.AutomaticControlDependencies() as c:
+
+        def true_fn():
+          return 0.0
+
+        def false_fn():
+          v.assign(v + 4)
+          return 1.0
+
+        cond.cond(p, true_fn, false_fn)
+        val = v.read_value()
+        val = c.mark_as_return(val)
+      self.assertAllEqual(val.eval(feed_dict={p: False}), 5.0)
+      self.assertAllEqual(val.eval(feed_dict={p: True}), 5.0)
+
+  @test_util.run_v1_only("b/120545219")
+  def testCondOneBranchUpdateBefore(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      p = array_ops.placeholder(dtype=dtypes.bool)
+      with acd.AutomaticControlDependencies() as c:
+        v.assign(v * 2)
+
+        def true_fn():
+          return 0.0
+
+        def false_fn():
+          v.assign(v + 4)
+          return 1.0
+
+        cond.cond(p, true_fn, false_fn)
+        val = v.read_value()
+        val = c.mark_as_return(val)
+      self.assertAllEqual(val.eval(feed_dict={p: False}), 6.0)
+      self.assertAllEqual(val.eval(feed_dict={p: True}), 12.0)
+
+  @test_util.run_v1_only("b/120545219")
+  def testCondOneBranchUpdateAfter(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      p = array_ops.placeholder(dtype=dtypes.bool)
+      with acd.AutomaticControlDependencies() as c:
+
+        def true_fn():
+          return 0.0
+
+        def false_fn():
+          v.assign(v + 4)
+          return 1.0
+
+        cond.cond(p, true_fn, false_fn)
+        v.assign(v * 2)
+        val = v.read_value()
+        val = c.mark_as_return(val)
+      self.assertAllEqual(val.eval(feed_dict={p: False}), 10.0)
+      self.assertAllEqual(val.eval(feed_dict={p: True}), 20.0)
+
+  def testFunctionWhileLoopWithCapturedLoopVars(self):
+    n = 3
+    x = constant_op.constant(list(range(n)))
+
+    @def_function.function
+    def loop():
+      c = lambda i, x: i < n
+      b = lambda i, x: (i + 1, x + 1)
+      i, out = while_loop.while_loop(c, b, (0, x))
+      return i, out
+
+    i, out = loop()
+    self.assertEqual(int(i), 3)
+    self.assertAllEqual(out, [3, 4, 5])
+
+  def testDecorator(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+
+      @acd.automatic_control_dependencies
+      def f():
+        v.assign(v + 1)
+        v.assign(2 * v)
+        return v.read_value()
+
+      self.assertAllEqual(f(), 4.0)
+
+  def testOptimizerInFunction(self):
+    def loss(v):
+      return v**2
+
+    optimizer = momentum.MomentumOptimizer(learning_rate=1.0, momentum=1.0)
+
+    @def_function.function
+    def train():
+      grad = backprop.implicit_grad(loss)(self.v)
+      optimizer.apply_gradients(grad)
+      return self.v.read_value()
+
+    self.v = resource_variable_ops.ResourceVariable(1.0)
+    value = train()
+    self.assertEqual(value.numpy(), -1.0)
+
+  def testReturningNonTensorRaisesError(self):
+    optimizer = momentum.MomentumOptimizer(learning_rate=1.0, momentum=1.0)
+    optimizer.apply_gradients = def_function.function(optimizer.apply_gradients)
+    v = resource_variable_ops.ResourceVariable(1.0)
+    grad = backprop.implicit_grad(lambda v: v**2)(v)
+
+    with self.assertRaisesRegex(TypeError,
+                                ".*must return zero or more Tensors.*"):
+      # TODO(akshayka): We might want to allow defun-ing Python functions
+      # that return operations (and just execute the op instead of running it).
+      optimizer.apply_gradients(grad)
+
+  # TODO(b/111663004): This should work when the outer context is graph
+  # building.
+  def testOptimizerNonSlotVarsInFunctionNoError(self):
+    def loss(v):
+      return v**2
+
+    optimizer = adam.AdamOptimizer(learning_rate=1.0)
+
+    @def_function.function
+    def train():
+      grad = backprop.implicit_grad(loss)(self.v)
+      optimizer.apply_gradients(grad)
+      return self.v.read_value()
+
+    self.v = resource_variable_ops.ResourceVariable(1.0)
+    train()
+
+  def testOptimizerInFunctionWithCapturedVariable(self):
+    v = resource_variable_ops.ResourceVariable(1.0)
+    def loss():
+      return v**2
+
+    optimizer = momentum.MomentumOptimizer(learning_rate=1.0, momentum=1.0)
+
+    @def_function.function
+    def train():
+      grad = backprop.implicit_grad(loss)()
+      optimizer.apply_gradients(grad)
+
+    train()
+    self.assertEqual(v.numpy(), -1.0)
+
+  def testRepeatedResourceInput(self):
+    var = resource_variable_ops.ResourceVariable(1.0)
+
+    @def_function.function
+    def inner(var1, var2):
+      return (resource_variable_ops.read_variable_op(var1, dtypes.float32) +
+              resource_variable_ops.read_variable_op(var2, dtypes.float32))
+
+    @def_function.function
+    def outer():
+      return inner(var.handle, var.handle)
+
+    self.assertEqual(self.evaluate(outer()), 2.0)
+
+  def testManualControlDepMonitoringAttrNotAdded(self):
+    with context.graph_mode(), self.cached_session():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      self.evaluate(variables.global_variables_initializer())
+      with acd.AutomaticControlDependencies():
+        read_op1 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        read_op2 = gen_resource_variable_ops.read_variable_op(
+            v.handle, v.dtype).op
+        assign_op = gen_resource_variable_ops.assign_variable_op(
+            v.handle, v + 1)
+      # Writes should have control deps automatically added from "all" reads
+      # since last write or start of the code block.
+      self.assertIn(read_op1, assign_op.control_inputs)
+      self.assertIn(read_op2, assign_op.control_inputs)
+      # But, we shouldn't add the monitoring attribute in this case.
+      with self.assertRaises(ValueError):
+        assign_op.get_attr("_has_manual_control_dependencies")
+      with self.assertRaises(ValueError):
+        read_op1.get_attr("_has_manual_control_dependencies")
+      with self.assertRaises(ValueError):
+        read_op2.get_attr("_has_manual_control_dependencies")
+
+  def testReadOnlyResourceInputsUnsorted(self):
+    from tensorflow.python.framework import auto_control_deps_utils as acd_utils
+    from unittest import mock
+
+    attr = acd_utils.READ_ONLY_RESOURCE_INPUTS_ATTR
+
+    def get_attr_returning(indices):
+      def _get_attr(name):
+        if name == attr:
+          return indices
+        raise ValueError()
+
+      return _get_attr
+
+    # 1. Unsorted inputs
+    input0 = mock.MagicMock()
+    input0.dtype = dtypes.resource
+    input1 = mock.MagicMock()
+    input1.dtype = dtypes.int32
+    input2 = mock.MagicMock()
+    input2.dtype = dtypes.resource
+
+    op = mock.MagicMock()
+    op.type = "SomeOp"
+    op.inputs = [input0, input1, input2]
+    op.get_attr.side_effect = get_attr_returning([2, 0])
+
+    # _get_read_only_resource_input_indices_op returns both indices, sorted.
+    self.assertEqual(
+        acd_utils._get_read_only_resource_input_indices_op(op), [0, 2]
+    )
+
+    # get_read_write_resource_inputs marks both as reads, none as writes.
+    reads, writes = acd_utils.get_read_write_resource_inputs(op)
+    self.assertIn(input0, reads)
+    self.assertIn(input2, reads)
+    self.assertEqual(len(writes), 0)
+
+    # 2. Empty indices
+    op2 = mock.MagicMock()
+    op2.type = "SomeOp"
+    op2.inputs = [input0, input1, input2]
+    op2.get_attr.side_effect = get_attr_returning([])
+
+    self.assertEqual(
+        acd_utils._get_read_only_resource_input_indices_op(op2), []
+    )
+    reads2, writes2 = acd_utils.get_read_write_resource_inputs(op2)
+    self.assertEqual(len(reads2), 0)
+    self.assertIn(input0, writes2)
+    self.assertIn(input2, writes2)
+
+    # 3. Mixed reads and writes (index 0 is read, index 2 is write)
+    op3 = mock.MagicMock()
+    op3.type = "SomeOp"
+    op3.inputs = [input0, input1, input2]
+    op3.get_attr.side_effect = get_attr_returning([0])
+
+    self.assertEqual(
+        acd_utils._get_read_only_resource_input_indices_op(op3), [0]
+    )
+    reads3, writes3 = acd_utils.get_read_write_resource_inputs(op3)
+    self.assertIn(input0, reads3)
+    self.assertNotIn(input2, reads3)
+    self.assertIn(input2, writes3)
+    self.assertNotIn(input0, writes3)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/framework/auto_control_deps_utils.py

@@ -0,0 +1,162 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for AutomaticControlDependencies."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.util import object_identity
+
+READ_ONLY_RESOURCE_INPUTS_ATTR = "_read_only_resource_inputs"
+RESOURCE_READ_OPS = set()
+
+
+COLLECTIVE_MANAGER_IDS = "_collective_manager_ids"
+
+
+def register_read_only_resource_op(op_type):
+  """Declares that `op_type` does not update its touched resource."""
+  RESOURCE_READ_OPS.add(op_type)
+
+
+def get_read_only_resource_input_indices_graph(func_graph):
+  """Returns sorted list of read-only resource indices in func_graph.inputs."""
+  result = []
+  # A cache to store the read only resource inputs of an Op.
+  # Operation -> ObjectIdentitySet of resource handles.
+  op_read_only_resource_inputs = {}
+  for input_index, t in enumerate(func_graph.inputs):
+    if t.dtype != dtypes.resource:
+      continue
+    read_only = True
+    for op in t.consumers():
+      if op in op_read_only_resource_inputs:
+        if t not in op_read_only_resource_inputs[op]:
+          read_only = False
+          break
+      else:
+        indices = _get_read_only_resource_input_indices_op(op)
+        op_read_only_resource_inputs[op] = object_identity.ObjectIdentitySet(
+            [op.inputs[i] for i in indices])
+        if t not in op_read_only_resource_inputs[op]:
+          read_only = False
+          break
+    if read_only:
+      result.append(input_index)
+  return result
+
+
+def _get_read_only_resource_input_indices_op(op):
+  """Returns sorted list of read-only resource indices in op.inputs."""
+  if op.type in RESOURCE_READ_OPS:
+    return [i for i, t in enumerate(op.inputs) if t.dtype == dtypes.resource]
+
+  try:
+    read_only_input_indices = op.get_attr(READ_ONLY_RESOURCE_INPUTS_ATTR)
+  except ValueError:
+    # Attr was not set. Add all resource inputs to `writes` and return.
+    return []
+
+  read_only_input_indices_set = set(read_only_input_indices)
+  result = []
+  for i, t in enumerate(op.inputs):
+    if t.dtype != dtypes.resource:
+      continue
+    if i in read_only_input_indices_set:
+      result.append(i)
+
+  return result
+
+
+def get_read_write_resource_inputs(op):
+  """Returns a tuple of resource reads, writes in op.inputs.
+
+  Args:
+    op: Operation
+
+  Returns:
+    A 2-tuple of ObjectIdentitySets, the first entry containing read-only
+    resource handles and the second containing read-write resource handles in
+    `op.inputs`.
+  """
+  reads = object_identity.ObjectIdentitySet()
+  writes = object_identity.ObjectIdentitySet()
+
+  if op.type in RESOURCE_READ_OPS:
+    # Add all resource inputs to `reads` and return.
+    reads.update(t for t in op.inputs if t.dtype == dtypes.resource)
+    return (reads, writes)
+
+  try:
+    read_only_input_indices = op.get_attr(READ_ONLY_RESOURCE_INPUTS_ATTR)
+  except ValueError:
+    # Attr was not set. Add all resource inputs to `writes` and return.
+    writes.update(t for t in op.inputs if t.dtype == dtypes.resource)
+    return (reads, writes)
+
+  read_only_input_indices_set = set(read_only_input_indices)
+  for i, t in enumerate(op.inputs):
+    if t.dtype != dtypes.resource:
+      continue
+    if i in read_only_input_indices_set:
+      reads.add(t)
+    else:
+      writes.add(t)
+  return (reads, writes)
+
+
+def _op_writes_to_resource(handle, op):
+  """Returns whether op writes to resource handle.
+
+  Args:
+    handle: Resource handle. Must be an input of `op`.
+    op: Operation.
+
+  Returns:
+    Returns False if op is a read-only op registered using
+    `register_read_only_resource_op` or if `handle` is an input at one of
+    the indices in the `READ_ONLY_RESOURCE_INPUTS_ATTR` attr of the op, True
+    otherwise.
+
+  Raises:
+    ValueError: if `handle` is not an input of `op`.
+  """
+  if op.type in RESOURCE_READ_OPS:
+    return False
+  input_index = _input_index(op, handle)
+  try:
+    read_only_input_indices = op.get_attr(READ_ONLY_RESOURCE_INPUTS_ATTR)
+  except ValueError:
+    # Attr was not set. Conservatively assume that the resource is written to.
+    return True
+  return input_index not in read_only_input_indices
+
+
+def _input_index(op, handle):
+  """Returns the index of `handle` in `op.inputs`.
+
+  Args:
+    op: Operation.
+    handle: Resource handle.
+
+  Returns:
+    Index in `op.inputs` receiving the resource `handle`.
+
+  Raises:
+    ValueError: If handle and its replicated input are both not found in
+    `op.inputs`.
+  """
+  for i, t in enumerate(op.inputs):
+    if handle is t:
+      return i
+  raise ValueError(f"{handle!s} not in list of inputs for op: {op!r}")

tensorflow/python/framework/byte_swap_tensor.py

@@ -0,0 +1,103 @@
+# 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.
+# ==============================================================================
+
+"""Utilities for byte swapping the tensor content."""
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python.framework import dtypes
+
+# Based on tensor_bundle/byte_swap.cc
+byte_swappable = [
+    dtypes.float16,
+    dtypes.float32,
+    dtypes.float64,
+    dtypes.bfloat16,
+    dtypes.complex64,
+    dtypes.complex128,
+    dtypes.uint16,
+    dtypes.uint32,
+    dtypes.uint64,
+    dtypes.int16,
+    dtypes.int32,
+    dtypes.int64,
+    dtypes.qint16,
+    dtypes.quint16,
+    dtypes.qint32,
+]
+
+
+def byte_swap_tensor_content(tensor, from_endiness, to_endiness):
+  """Byte swaps.
+
+  Args:
+    tensor: Target tensor to change endiness.
+    from_endiness: The original endianness format. "big" or "little"
+    to_endiness: The target endianness format. "big" or "little"
+  """
+  if tensor.dtype in byte_swappable:
+    tshape = tensor.tensor_shape.dim
+    tensor_bytes = tensor.tensor_content
+    if tensor_bytes:
+      tensor_size = 1
+      for sz in tshape:
+        if sz.size != 0:
+          tensor_size *= sz.size
+      chunksize = len(tensor_bytes) // tensor_size
+      # Split tensor_data into chunks for byte swapping.
+      to_swap = [
+          tensor_bytes[i : i + chunksize]
+          for i in range(0, len(tensor_bytes), chunksize)
+      ]
+      # Swap and replace tensor_content.
+      tensor.tensor_content = b"".join(
+          [
+              int.from_bytes(byteswap, from_endiness).to_bytes(
+                  chunksize, to_endiness
+              )
+              for byteswap in to_swap
+          ]
+      )
+
+
+def swap_tensor_content_in_graph_function(
+    graph_def, from_endiness, to_endiness
+):
+  """Fix endiness of tensor contents.
+
+  Args:
+    graph_def: Target graph_def to change endiness.
+    from_endiness: The original endianness format. "big" or "little"
+    to_endiness: The target endianness format. "big" or "little"
+  """
+  if isinstance(graph_def, meta_graph_pb2.MetaGraphDef):
+    functions = graph_def.graph_def.library.function
+  elif isinstance(graph_def, graph_pb2.GraphDef):
+    functions = graph_def.library.function
+  else:
+    return
+  for function in functions:
+    node_def = function.node_def
+    for node in node_def:
+      if node.op == "Const":
+        tensor = node.attr["value"].tensor
+        byte_swap_tensor_content(tensor, from_endiness, to_endiness)
+
+
+def swap_tensor_content_in_graph_node(graph_def, from_endiness, to_endiness):
+  for node in graph_def.node:
+    if node.op == "Const":
+      tensor = node.attr["value"].tensor
+      byte_swap_tensor_content(tensor, from_endiness, to_endiness)

tensorflow/python/framework/c_api_util.py

@@ -0,0 +1,239 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utilities for using the TensorFlow C API."""
+
+import contextlib
+from tensorflow.core.framework import api_def_pb2
+from tensorflow.core.framework import op_def_pb2
+from tensorflow.python.client import pywrap_tf_session as c_api
+from tensorflow.python.util import compat
+from tensorflow.python.util import tf_contextlib
+
+
+class AlreadyGarbageCollectedError(Exception):
+
+  def __init__(self, name, obj_type):
+    super(AlreadyGarbageCollectedError,
+          self).__init__(f"{name} of type {obj_type} has already been garbage "
+                         f"collected and cannot be called.")
+
+
+# FIXME(b/235488206): Convert all Scoped objects to the context manager
+# to protect against deletion during use when the object is attached to
+# an attribute.
+class UniquePtr(object):
+  """Wrapper around single-ownership C-API objects that handles deletion."""
+
+  __slots__ = ["_obj", "deleter", "name", "type_name"]
+
+  def __init__(self, name, obj, deleter):
+    # '_' prefix marks _obj private, but unclear if it is required also to
+    # maintain a special CPython destruction order.
+    self._obj = obj
+    self.name = name
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we may have already deleted other modules. By capturing the
+    # DeleteGraph function here, we retain the ability to cleanly destroy the
+    # graph at shutdown, which satisfies leak checkers.
+    self.deleter = deleter
+    self.type_name = str(type(obj))
+
+  @contextlib.contextmanager
+  def get(self):
+    """Yields the managed C-API Object, guaranteeing aliveness.
+
+    This is a context manager. Inside the context the C-API object is
+    guaranteed to be alive.
+
+    Raises:
+      AlreadyGarbageCollectedError: if the object is already deleted.
+    """
+    # Thread-safety: self.__del__ never runs during the call of this function
+    # because there is a reference to self from the argument list.
+    if self._obj is None:
+      raise AlreadyGarbageCollectedError(self.name, self.type_name)
+    yield self._obj
+
+  def __del__(self):
+    obj = self._obj
+    if obj is not None:
+      self._obj = None
+      self.deleter(obj)
+
+
+class ScopedTFStatus(object):
+  """Wrapper around TF_Status that handles deletion."""
+
+  __slots__ = ["status"]
+
+  def __init__(self):
+    self.status = c_api.TF_NewStatus()
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteStatus is not None:
+      c_api.TF_DeleteStatus(self.status)
+
+
+class ScopedTFImportGraphDefOptions(object):
+  """Wrapper around TF_ImportGraphDefOptions that handles deletion."""
+
+  __slots__ = ["options"]
+
+  def __init__(self):
+    self.options = c_api.TF_NewImportGraphDefOptions()
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteImportGraphDefOptions is not None:
+      c_api.TF_DeleteImportGraphDefOptions(self.options)
+
+
+class ScopedTFImportGraphDefResults(object):
+  """Wrapper around TF_ImportGraphDefResults that handles deletion."""
+
+  __slots__ = ["results"]
+
+  def __init__(self, results):
+    self.results = results
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteImportGraphDefResults is not None:
+      c_api.TF_DeleteImportGraphDefResults(self.results)
+
+
+class ScopedTFFunction(UniquePtr):
+  """Wrapper around TF_Function that handles deletion."""
+
+  def __init__(self, func, name):
+    super(ScopedTFFunction, self).__init__(
+        name=name, obj=func, deleter=c_api.TF_DeleteFunction)
+
+
+class ScopedTFBuffer(object):
+  """An internal class to help manage the TF_Buffer lifetime."""
+
+  __slots__ = ["buffer"]
+
+  def __init__(self, buf_string):
+    self.buffer = c_api.TF_NewBufferFromString(compat.as_bytes(buf_string))
+
+  def __del__(self):
+    c_api.TF_DeleteBuffer(self.buffer)
+
+
+class ApiDefMap(object):
+  """Wrapper around TF_ApiDefMap that handles querying and deletion.
+
+  The OpDef protos are also stored in this class so that they could
+  be queried by op name.
+  """
+
+  __slots__ = ["_api_def_map", "_op_per_name"]
+
+  def __init__(self):
+    op_def_proto = op_def_pb2.OpList()
+    buf = c_api.TF_GetAllOpList()
+    try:
+      op_def_proto.ParseFromString(c_api.TF_GetBuffer(buf))
+      self._api_def_map = c_api.TF_NewApiDefMap(buf)
+    finally:
+      c_api.TF_DeleteBuffer(buf)
+
+    self._op_per_name = {}
+    for op in op_def_proto.op:
+      self._op_per_name[op.name] = op
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteApiDefMap is not None:
+      c_api.TF_DeleteApiDefMap(self._api_def_map)
+
+  def put_api_def(self, text):
+    c_api.TF_ApiDefMapPut(self._api_def_map, text, len(text))
+
+  def get_api_def(self, op_name):
+    api_def_proto = api_def_pb2.ApiDef()
+    buf = c_api.TF_ApiDefMapGet(self._api_def_map, op_name, len(op_name))
+    if buf is None:
+      raise ValueError(f"No api_def found for op name {op_name}.")
+    try:
+      api_def_proto.ParseFromString(c_api.TF_GetBuffer(buf))
+    finally:
+      c_api.TF_DeleteBuffer(buf)
+    return api_def_proto
+
+  def get_op_def(self, op_name):
+    if op_name in self._op_per_name:
+      return self._op_per_name[op_name]
+    raise ValueError(f"No op_def found for op name {op_name}.")
+
+  def op_names(self):
+    return self._op_per_name.keys()
+
+
+@tf_contextlib.contextmanager
+def tf_buffer(data=None):
+  """Context manager that creates and deletes TF_Buffer.
+
+  Example usage:
+    with tf_buffer() as buf:
+      # get serialized graph def into buf
+      ...
+      proto_data = c_api.TF_GetBuffer(buf)
+      graph_def.ParseFromString(compat.as_bytes(proto_data))
+    # buf has been deleted
+
+    with tf_buffer(some_string) as buf:
+      c_api.TF_SomeFunction(buf)
+    # buf has been deleted
+
+  Args:
+    data: An optional `bytes`, `str`, or `unicode` object. If not None, the
+      yielded buffer will contain this data.
+
+  Yields:
+    Created TF_Buffer
+  """
+  if data:
+    buf = c_api.TF_NewBufferFromString(compat.as_bytes(data))
+  else:
+    buf = c_api.TF_NewBuffer()
+  try:
+    yield buf
+  finally:
+    c_api.TF_DeleteBuffer(buf)
+
+
+def tf_output(c_op, index):
+  """Returns a wrapped TF_Output with specified operation and index.
+
+  Args:
+    c_op: wrapped TF_Operation
+    index: integer
+
+  Returns:
+    Wrapped TF_Output
+  """
+  ret = c_api.TF_Output()
+  ret.oper = c_op
+  ret.index = index
+  return ret

tensorflow/python/framework/c_api_util_test.py

@@ -0,0 +1,126 @@
+# Copyright 2016 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 c_api utils."""
+import gc
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+
+class ApiDefMapTest(test_util.TensorFlowTestCase):
+
+  def testApiDefMapOpNames(self):
+    api_def_map = c_api_util.ApiDefMap()
+    self.assertIn("Add", api_def_map.op_names())
+
+  def testApiDefMapGet(self):
+    api_def_map = c_api_util.ApiDefMap()
+    op_def = api_def_map.get_op_def("Add")
+    self.assertEqual(op_def.name, "Add")
+    api_def = api_def_map.get_api_def("Add")
+    self.assertEqual(api_def.graph_op_name, "Add")
+
+  def testApiDefMapGetInvalidOp(self):
+    api_def_map = c_api_util.ApiDefMap()
+    with self.assertRaises(ValueError):
+      api_def_map.get_api_def("InvalidOperationName")
+
+  def testApiDefMapPutThenGet(self):
+    api_def_map = c_api_util.ApiDefMap()
+    api_def_text = """
+op {
+  graph_op_name: "Add"
+  summary: "Returns x + y element-wise."
+  description: <<END
+*NOTE*: `Add` supports broadcasting. `AddN` does not. More about broadcasting
+[here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+END
+}
+"""
+    api_def_map.put_api_def(api_def_text)
+    api_def = api_def_map.get_api_def("Add")
+    self.assertEqual(api_def.graph_op_name, "Add")
+    self.assertEqual(api_def.summary, "Returns x + y element-wise.")
+
+
+class UniquePtrTest(test_util.TensorFlowTestCase):
+
+  def setUp(self):
+
+    super(UniquePtrTest, self).setUp()
+
+    class MockClass:
+
+      def __init__(self):
+        self.deleted = False
+
+    def deleter(obj):
+      obj.deleted = True
+
+    self.obj = MockClass()
+    self.deleter = deleter
+
+  def testLifeCycle(self):
+    self.assertFalse(self.obj.deleted)
+
+    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
+
+    with a.get() as obj:
+      self.assertIs(obj, self.obj)
+
+    del a
+    gc.collect()
+    self.assertTrue(self.obj.deleted)
+
+  def testSafeUnderRaceCondition(self):
+    self.assertFalse(self.obj.deleted)
+
+    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
+
+    with a.get() as obj:
+      self.assertIs(obj, self.obj)
+      # The del below mimics a potential race condition.
+      # 'a' could be owned by a different thread, and this thread not
+      # necessarily hold a long-term reference to a.
+      del a
+      gc.collect()
+      self.assertFalse(obj.deleted)
+
+    gc.collect()
+    self.assertTrue(self.obj.deleted)
+
+  def testRaiseAfterDeleted(self):
+    self.assertFalse(self.obj.deleted)
+
+    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
+
+    # The __del__ below mimics a partially started deletion, potentially
+    # started from another thread.
+    # 'a' could be owned by a different thread, and this thread not
+    # necessarily hold a long-term reference to a.
+    a.__del__()
+    self.assertTrue(self.obj.deleted)
+
+    with self.assertRaisesRegex(c_api_util.AlreadyGarbageCollectedError,
+                                "MockClass"):
+      with a.get():
+        pass
+
+    gc.collect()
+    self.assertTrue(self.obj.deleted)
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/framework/combinations.py

@@ -0,0 +1,83 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""This module customizes `test_combinations` for Tensorflow.
+
+Additionally it provides `generate()`, `combine()` and `times()` with Tensorflow
+customizations as a default.
+"""
+
+import functools
+
+from tensorflow.python import tf2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_combinations
+from tensorflow.python.util.tf_export import tf_export
+
+
+class EagerGraphCombination(test_combinations.TestCombination):
+  """Run the test in Graph or Eager mode.
+
+  The optional `mode` parameter controls the test's execution mode.  Its
+  accepted values are "graph" or "eager" literals.
+  """
+
+  def context_managers(self, kwargs):
+    mode = kwargs.pop("mode", None)
+    if mode is None:
+      return []
+    elif mode == "eager":
+      return [context.eager_mode()]
+    elif mode == "graph":
+      return [ops.Graph().as_default(), context.graph_mode()]
+    else:
+      raise ValueError(
+          "Argument 'mode' must be either 'eager' or 'graph'. "
+          f"Received: {mode}.")
+
+  def parameter_modifiers(self):
+    return [test_combinations.OptionalParameter("mode")]
+
+
+class TFVersionCombination(test_combinations.TestCombination):
+  """Control the execution of the test in TF1.x and TF2.
+
+  If TF2 is enabled then a test with TF1 test is going to be skipped and vice
+  versa.
+
+  Test targets continuously run in TF2 thanks to the tensorflow.v2 TAP target.
+  A test can be run in TF2 with bazel by passing --test_env=TF2_BEHAVIOR=1.
+  """
+
+  def should_execute_combination(self, kwargs):
+    tf_api_version = kwargs.pop("tf_api_version", None)
+    if tf_api_version == 1 and tf2.enabled():
+      return (False, "Skipping a TF1.x test when TF2 is enabled.")
+    elif tf_api_version == 2 and not tf2.enabled():
+      return (False, "Skipping a TF2 test when TF2 is not enabled.")
+    return (True, None)
+
+  def parameter_modifiers(self):
+    return [test_combinations.OptionalParameter("tf_api_version")]
+
+
+generate = functools.partial(
+    test_combinations.generate,
+    test_combinations=(EagerGraphCombination(), TFVersionCombination()))
+combine = test_combinations.combine
+times = test_combinations.times
+NamedObject = test_combinations.NamedObject
+
+tf_export("__internal__.test.combinations.generate", v1=[])(generate)

tensorflow/python/framework/common_shapes.py

@@ -0,0 +1,108 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A library of common shape functions."""
+import itertools
+
+from tensorflow.python.framework import tensor_shape
+
+
+def _broadcast_shape_helper(shape_x, shape_y):
+  """Helper functions for is_broadcast_compatible and broadcast_shape.
+
+  Args:
+    shape_x: A `TensorShape`
+    shape_y: A `TensorShape`
+
+  Returns:
+    Returns None if the shapes are not broadcast compatible,
+    a list of the broadcast dimensions otherwise.
+  """
+  # To compute the broadcasted dimensions, we zip together shape_x and shape_y,
+  # and pad with 1 to make them the same length.
+  broadcasted_dims = reversed(
+      list(
+          itertools.zip_longest(
+              reversed(shape_x.dims),
+              reversed(shape_y.dims),
+              fillvalue=tensor_shape.Dimension(1))))
+  # Next we combine the dimensions according to the numpy broadcasting rules.
+  # http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
+  return_dims = []
+  for (dim_x, dim_y) in broadcasted_dims:
+    if dim_x.value is None or dim_y.value is None:
+      # One or both dimensions is unknown. If either dimension is greater than
+      # 1, we assume that the program is correct, and the other dimension will
+      # be broadcast to match it.
+      # TODO(mrry): If we eliminate the shape checks in C++, we must still
+      # assert that the unknown dim is either 1 or the same as the known dim.
+      if dim_x.value is not None and dim_x.value > 1:
+        return_dims.append(dim_x)
+      elif dim_y.value is not None and dim_y.value > 1:
+        return_dims.append(dim_y)
+      else:
+        return_dims.append(None)
+    elif dim_x.value == 1:
+      # We will broadcast dim_x to dim_y.
+      return_dims.append(dim_y)
+    elif dim_y.value == 1:
+      # We will broadcast dim_y to dim_x.
+      return_dims.append(dim_x)
+    elif dim_x.value == dim_y.value:
+      # The dimensions are compatible, so output is the same size in that
+      # dimension.
+      return_dims.append(dim_x.merge_with(dim_y))
+    else:
+      return None
+  return return_dims
+
+
+def is_broadcast_compatible(shape_x, shape_y):
+  """Returns True if `shape_x` and `shape_y` are broadcast compatible.
+
+  Args:
+    shape_x: A `TensorShape`
+    shape_y: A `TensorShape`
+
+  Returns:
+    True if a shape exists that both `shape_x` and `shape_y` can be broadcasted
+    to.  False otherwise.
+  """
+  if shape_x.ndims is None or shape_y.ndims is None:
+    return False
+  return _broadcast_shape_helper(shape_x, shape_y) is not None
+
+
+def broadcast_shape(shape_x, shape_y):
+  """Returns the broadcasted shape between `shape_x` and `shape_y`.
+
+  Args:
+    shape_x: A `TensorShape`
+    shape_y: A `TensorShape`
+
+  Returns:
+    A `TensorShape` representing the broadcasted shape.
+
+  Raises:
+    ValueError: If the two shapes can not be broadcasted.
+  """
+  if shape_x.ndims is None or shape_y.ndims is None:
+    return tensor_shape.unknown_shape()
+  return_dims = _broadcast_shape_helper(shape_x, shape_y)
+  if return_dims is None:
+    raise ValueError('Incompatible shapes for broadcasting. Two shapes are '
+                     'compatible if for each dimension pair they are either '
+                     'equal or one of them is 1. '
+                     f'Received: {shape_x} and {shape_y}.')
+  return tensor_shape.TensorShape(return_dims)

tensorflow/python/framework/common_shapes_test.py

@@ -0,0 +1,205 @@
+# Copyright 2016 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 common shapes."""
+import numpy as np
+
+from tensorflow.python.framework import common_shapes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+
+class CommonShapesTest(test_util.TensorFlowTestCase):
+
+  # Asserts that we get the same result with numpy (for known shapes), and that
+  # the order of arguments does not matter (i.e., broadcasting is reflexive).
+  def _assert_incompatible_broadcast(self, shape1, shape2):
+    if shape1.dims is not None and shape2.dims is not None:
+      zeros1 = np.zeros(shape1.as_list())
+      zeros2 = np.zeros(shape2.as_list())
+      with self.assertRaises(ValueError):
+        np.broadcast(zeros1, zeros2)
+      with self.assertRaises(ValueError):
+        np.broadcast(zeros2, zeros1)
+    self.assertFalse(common_shapes.is_broadcast_compatible(shape1, shape2))
+    self.assertFalse(common_shapes.is_broadcast_compatible(shape2, shape1))
+    with self.assertRaises(ValueError):
+      common_shapes.broadcast_shape(shape1, shape2)
+    with self.assertRaises(ValueError):
+      common_shapes.broadcast_shape(shape2, shape1)
+
+  # Asserts that we get the same result with numpy (for known shapes), and that
+  # the order of arguments does not matter (i.e., broadcasting is reflexive).
+  def _assert_broadcast(self, expected, shape1, shape2):
+    if shape1.dims is not None and shape2.dims is not None:
+      expected_np = expected.as_list()
+      zeros1 = np.zeros(shape1.as_list())
+      zeros2 = np.zeros(shape2.as_list())
+      self.assertAllEqual(expected_np, np.broadcast(zeros1, zeros2).shape)
+      self.assertAllEqual(expected_np, np.broadcast(zeros2, zeros1).shape)
+      self.assertEqual(
+          expected, common_shapes.broadcast_shape(shape1, shape2))
+      self.assertEqual(
+          expected, common_shapes.broadcast_shape(shape2, shape1))
+    else:
+      self.assertEqual(expected, common_shapes.broadcast_shape(shape1, shape2))
+      self.assertEqual(expected, common_shapes.broadcast_shape(shape2, shape1))
+
+  def testBroadcast_one_dimension(self):
+    s1 = tensor_shape.TensorShape([5])
+    s2 = tensor_shape.TensorShape([7])
+
+    unknown = tensor_shape.unknown_shape()
+    scalar = tensor_shape.TensorShape([])
+    expanded_scalar = tensor_shape.TensorShape([1])
+
+    # Tensors with same shape should have the same broadcast result.
+    for shape in (s1, s2, unknown, scalar, expanded_scalar):
+      self._assert_broadcast(expected=shape, shape1=shape, shape2=shape)
+
+    # [] and [1] act like identity.
+    self._assert_broadcast(expected=s1, shape1=s1, shape2=scalar)
+    self._assert_broadcast(expected=s2, shape1=s2, shape2=scalar)
+    self._assert_broadcast(expected=s1, shape1=s1, shape2=expanded_scalar)
+    self._assert_broadcast(expected=s2, shape1=s2, shape2=expanded_scalar)
+
+    self._assert_broadcast(expected=unknown, shape1=s1, shape2=unknown)
+    self._assert_broadcast(expected=unknown, shape1=s2, shape2=unknown)
+
+    self._assert_broadcast(
+        expected=expanded_scalar, shape1=scalar, shape2=expanded_scalar)
+
+    self._assert_incompatible_broadcast(shape1=s1, shape2=s2)
+
+  def testBroadcast_many_dimensions(self):
+    unknown = tensor_shape.unknown_shape()
+    shape_0 = tensor_shape.TensorShape([])
+    shape_1 = tensor_shape.TensorShape([1])
+    shape_4 = tensor_shape.TensorShape([4])
+    shape_1x4 = tensor_shape.TensorShape([1, 4])
+    shape_4x1 = tensor_shape.TensorShape([4, 1])
+    shape_3x4 = tensor_shape.TensorShape([3, 4])
+    shape_4x3 = tensor_shape.TensorShape([4, 3])
+
+    # Tensors with same shape should have the same broadcast result.
+    for shape in (
+        shape_0, shape_1, shape_4, shape_1x4, shape_4x1, shape_3x4, shape_4x3):
+      self._assert_broadcast(expected=shape, shape1=shape, shape2=shape)
+
+    # [] and [1] act like identity.
+    for identity in (shape_0, shape_1):
+      for shape in (shape_4, shape_1x4, shape_4x1, shape_3x4, shape_4x3):
+        self._assert_broadcast(expected=shape, shape1=identity, shape2=shape)
+
+    # Unknown in, unknown out.
+    for shape in (shape_4, shape_1x4, shape_4x1, shape_3x4, shape_4x3):
+      self._assert_broadcast(expected=unknown, shape1=shape, shape2=unknown)
+
+    self._assert_broadcast(expected=shape_1x4, shape1=shape_4, shape2=shape_1x4)
+    shape_4x4 = tensor_shape.TensorShape([4, 4])
+    self._assert_broadcast(expected=shape_4x4, shape1=shape_4, shape2=shape_4x1)
+    self._assert_broadcast(expected=shape_3x4, shape1=shape_4, shape2=shape_3x4)
+    self._assert_incompatible_broadcast(shape1=shape_4, shape2=shape_4x3)
+    self._assert_broadcast(
+        expected=shape_4x4, shape1=shape_1x4, shape2=shape_4x1)
+    self._assert_broadcast(
+        expected=shape_3x4, shape1=shape_1x4, shape2=shape_3x4)
+    self._assert_incompatible_broadcast(shape1=shape_1x4, shape2=shape_4x3)
+    self._assert_incompatible_broadcast(shape1=shape_4x1, shape2=shape_3x4)
+    self._assert_broadcast(
+        expected=shape_4x3, shape1=shape_4x1, shape2=shape_4x3)
+    self._assert_incompatible_broadcast(shape1=shape_3x4, shape2=shape_4x3)
+
+  # Asserts that the order of arguments does not matter (i.e., broadcasting is
+  # reflexive).
+  def _assert_broadcast_with_unknown_dims(self, expected, shape1, shape2):
+    actual_dims = common_shapes.broadcast_shape(shape1, shape2).dims
+    reflexive_actual_dims = common_shapes.broadcast_shape(shape2, shape1).dims
+
+    if actual_dims is None:
+      self.assertIsNone(reflexive_actual_dims)
+    elif reflexive_actual_dims is None:
+      self.assertIsNone(actual_dims)
+    else:
+      self.assertEqual(len(actual_dims), len(reflexive_actual_dims))
+      for actual_dim, reflexive_actual_dim in zip(
+          actual_dims, reflexive_actual_dims):
+        self.assertEqual(actual_dim.value, reflexive_actual_dim.value)
+
+    expected_dims = expected.dims
+    if expected_dims is None:
+      self.assertIsNone(actual_dims)
+    elif actual_dims is None:
+      self.assertIsNone(expected_dims)
+    else:
+      self.assertEqual(len(expected_dims), len(actual_dims))
+      for expected_dim, actual_dim in zip(expected_dims, actual_dims):
+        self.assertEqual(expected_dim.value, actual_dim.value)
+
+  def testBroadcast_unknown_dims(self):
+    unknown = tensor_shape.unknown_shape()
+    shape_0 = tensor_shape.TensorShape([])
+    shape_1 = tensor_shape.TensorShape([1])
+    # pylint: disable=invalid-name
+    shape_U = tensor_shape.TensorShape([None])
+    shape_1xU = tensor_shape.TensorShape([1, None])
+    shape_Ux1 = tensor_shape.TensorShape([None, 1])
+    shape_4xU = tensor_shape.TensorShape([4, None])
+    shape_Ux4 = tensor_shape.TensorShape([None, 4])
+    # pylint: enable=invalid-name
+
+    # Tensors with same shape should have the same broadcast result.
+    for shape in (shape_U, shape_1xU, shape_Ux1, shape_4xU, shape_Ux4):
+      self._assert_broadcast_with_unknown_dims(
+          expected=shape, shape1=shape, shape2=shape)
+
+    # [] and [1] act like identity.
+    for identity in (shape_0, shape_1):
+      for shape in (shape_U, shape_1xU, shape_Ux1, shape_4xU, shape_Ux4):
+        self._assert_broadcast_with_unknown_dims(
+            expected=shape, shape1=identity, shape2=shape)
+
+    # Unknown in, unknown out.
+    for shape in (shape_U, shape_1xU, shape_Ux1, shape_4xU, shape_Ux4):
+      self._assert_broadcast_with_unknown_dims(
+          expected=unknown, shape1=shape, shape2=unknown)
+
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_1xU, shape1=shape_U, shape2=shape_1xU)
+    shape_UxU = tensor_shape.TensorShape([None, None])  # pylint: disable=invalid-name
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_UxU, shape1=shape_U, shape2=shape_Ux1)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_4xU, shape1=shape_U, shape2=shape_4xU)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_Ux4, shape1=shape_U, shape2=shape_Ux4)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_UxU, shape1=shape_1xU, shape2=shape_Ux1)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_4xU, shape1=shape_1xU, shape2=shape_4xU)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_Ux4, shape1=shape_1xU, shape2=shape_Ux4)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_4xU, shape1=shape_Ux1, shape2=shape_4xU)
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_Ux4, shape1=shape_Ux1, shape2=shape_Ux4)
+    shape_4x4 = tensor_shape.TensorShape([4, 4])
+    self._assert_broadcast_with_unknown_dims(
+        expected=shape_4x4, shape1=shape_4xU, shape2=shape_Ux4)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/framework/composite_tensor.py

@@ -0,0 +1,140 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tensor-like objects that are composed from tf.Tensors."""
+
+import abc
+
+from tensorflow.python import pywrap_tensorflow  # pylint: disable=unused-import
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__internal__.CompositeTensor", v1=[])
+class CompositeTensor(metaclass=abc.ABCMeta):
+  """Abstract base class for Tensor-like objects that are composed from Tensors.
+
+  Each `CompositeTensor` can be decomposed into a structured collection of
+  component `tf.Tensor`s, and reconstructed from those components.
+
+  The `tensorflow.python.util.nest` module has support for treating composite
+  tensors as structure, which makes it easy to flatten and reconstruct
+  composite tensors (or larger structures that contain composite tensors).
+  E.g.:
+
+  ```python
+  ct = ...  # Create a composite tensor.
+  flat_list_of_tensors = nest.flatten(ct, expand_composites=True)
+  transformed_list_of_tensors = ...  # do something with the flat tensors.
+  result = nest.pack_sequence_as(ct, transformed_list_of_tensors,
+                                 expand_composites=True)
+  ```
+  """
+
+  @abc.abstractproperty
+  def _type_spec(self):
+    """A `TypeSpec` describing the type of this value."""
+    raise NotImplementedError(f"{type(self).__name__}._type_spec()")
+
+  def _shape_invariant_to_type_spec(self, shape):
+    """Returns a TypeSpec given a shape invariant (used by `tf.while_loop`).
+
+    Args:
+      shape: A `tf.TensorShape` object.  The shape invariant for this
+        `CompositeTensor`, or `None` if a default shape invariant should be used
+        (based on the value of this `CompositeTensor`).
+
+    Returns:
+      A nested structure whose values are `tf.TensorShape` objects, specifying
+      the shape invariants for the tensors that comprise this `CompositeTensor`.
+    """
+    # New TypeSpec subclasses generally do not need to implement this --
+    # this method is used for backwards compatibility.  Users of tf.while_loop
+    # can specify a type by passing in TypeSpec instead.
+    raise NotImplementedError(
+        f"{type(self).__name__}._shape_invariant_to_type_spec")
+
+  def _consumers(self):
+    """Returns a list of `Operation`s that consume this `CompositeTensor`.
+
+    Returns:
+      A list of `Operation`s.
+
+    Raises:
+      RuntimeError: If this method is called while executing eagerly.
+    """
+    consumers = nest.flatten([
+        component.consumers()
+        for component in nest.flatten(self, expand_composites=True)
+        if getattr(component, "graph", None) is not None
+    ])
+    return list(set(consumers))
+
+  def __tf_tracing_type__(self, context):
+    return self._type_spec.__tf_tracing_type__(context)
+
+  def _convert_variables_to_tensors(self):
+    """Converts ResourceVariable components to Tensors.
+
+    Override this method to explicitly convert ResourceVariables embedded in the
+    CompositeTensor to Tensors. By default, it returns the CompositeTensor
+    unchanged.
+
+    Returns:
+      A CompositeTensor with all its ResourceVariable components converted to
+      Tensors.
+    """
+    return self
+
+
+def replace_composites_with_components(structure):
+  """Recursively replaces CompositeTensors with their components.
+
+  Args:
+    structure: A `nest`-compatible structure, possibly containing composite
+      tensors.
+
+  Returns:
+    A copy of `structure`, where each composite tensor has been replaced by
+    its components.  The result will contain no composite tensors.
+    Note that `nest.flatten(replace_composites_with_components(structure))`
+    returns the same value as `nest.flatten(structure)`.
+  """
+  if isinstance(structure, CompositeTensor):
+    return replace_composites_with_components(
+        structure._type_spec._to_components(structure))  # pylint: disable=protected-access
+  elif not nest.is_nested(structure):
+    return structure
+  else:
+    return nest.map_structure(
+        replace_composites_with_components, structure, expand_composites=False)
+
+
+def convert_variables_to_tensors(composite_tensor):
+  return composite_tensor._convert_variables_to_tensors()  # pylint: disable=protected-access
+
+
+# @TODO(edloper): Can we replace convert_to_tensor_or_xyz with just
+# convert_to_tensor_or_composite?  Alternatively, should composite tensors
+# register a dispatch override for tf.convert_to_tensor?
+
+# Note about the internal encoding of composite tensors when they are "lowered"
+# from Python objects to tensors. The usual encoding is "component encoding"
+# which uses the dense tensors that represent a composite tensor.
+# A second encoding, "batchable tensor list encoding", is used by datasets
+# and map_fn which in addition to supporting batching also can use ops
+# for encoding and decoding, e.g. for encoding/decoding to/from a
+# single variant that represents a composite tensor. Some internal properties
+# for type specs for composite tensors use `flat` as a nickname for
+# "batchable tensor list encoding". (e.g. `flat_tensor_specs`).

tensorflow/python/framework/composite_tensor_gradient.py

@@ -0,0 +1,185 @@
+# 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.
+# ==============================================================================
+"""Gradient support for Composite Tensors."""
+
+import abc
+import sys
+
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.util import nest
+
+
+# pylint:disable=g-import-not-at-top
+if sys.version_info >= (3, 8):
+  from typing import Protocol
+  from typing import runtime_checkable
+else:
+  from typing_extensions import Protocol
+  from typing_extensions import runtime_checkable
+# pylint:enable=g-import-not-at-top
+
+
+# TODO(xjun): Add CompositeTensorGradient support for SparseTensor,
+# StructuredTensor, and MaskedTensor.
+class CompositeTensorGradient(object, metaclass=abc.ABCMeta):
+  """Class used to help compute gradients for CompositeTensors.
+
+  This abstract base class defines two methods: `get_gradient_components`, which
+  returns the components of a value that should be included in gradients; and
+  `replace_gradient_components`, which replaces the gradient components in a
+  value.  These methods can be used to compute the gradient of a `y` with
+  respect to `x` (`grad(y, x)`) as follows:
+
+  * If `y` is a `CompositeTensor` with `CompositeTensorGradient` `cg` =
+    `y.__composite_gradient__`, then `grad(y, x)` =
+    `grad(cg.get_gradient_components(y), x)`.
+
+  * If `x` is a `CompositeTensor` with `CompositeTensorGradient` `cg` =
+    'x.__composite_gradient__', then `grad(y, x)` =
+    `cg.replace_gradient_components(x, grad(y, cg.get_gradient_components(x))`.
+  """
+
+  @abc.abstractmethod
+  def get_gradient_components(self, value):
+    """Returns the components of `value` that should be included in gradients.
+
+    This method may not call TensorFlow ops, since any new ops added to the
+    graph would not be properly tracked by the gradient mechanisms.
+
+    Args:
+      value: A `CompositeTensor` value.
+
+    Returns:
+      A nested structure of `Tensor` or `IndexedSlices`.
+    """
+    raise NotImplementedError(
+        f"{type(self).__name__}.get_gradient_components()")
+
+  @abc.abstractmethod
+  def replace_gradient_components(self, value, component_grads):
+    """Replaces the gradient components in `value` with `component_grads`.
+
+    Args:
+      value: A value with its gradient components compatible with
+        `component_grads`.
+      component_grads: A nested structure of `Tensor` or `IndexedSlices` or
+        `None` (for unconnected gradients).
+
+    Returns:
+      A copy of `value`, where the components that should be included in
+      gradients have been replaced by `component_grads`; or `None` (if
+      `component_grads` includes `None`).
+    """
+    raise NotImplementedError(
+        f"{type(self).__name__}.replace_gradient_components()")
+
+
+@runtime_checkable
+class CompositeTensorGradientProtocol(Protocol):
+  """Protocol for adding gradient support to CompositeTensors."""
+  __composite_gradient__: CompositeTensorGradient
+
+
+class WithValuesCompositeTensorGradient(CompositeTensorGradient):
+  """CompositeTensorGradient based on `T.values` and `T.with_values`."""
+
+  def get_gradient_components(self, value):
+    return value.values
+
+  def replace_gradient_components(self, value, component_grads):
+    return value.with_values(component_grads)
+
+
+def _get_tensors_for_gradient(x):
+  """Returns the Tensors in `x` that should be differentiated.
+
+  Args:
+    x: A `Tensor` or `CompositeTensor`.
+
+  Returns:
+    A `Tensor` or a nested structure of `Tensor`.
+  """
+  if not isinstance(x, composite_tensor.CompositeTensor):
+    return x
+
+  if not isinstance(x, CompositeTensorGradientProtocol):
+    raise ValueError(
+        f"Type {type(x).__name__} is not supported as a gradient source or "
+        "gradient target.")
+  composite_gradient = x.__composite_gradient__
+  gradient_components = composite_gradient.get_gradient_components(x)
+  if gradient_components is x:
+    return x
+  return nest.map_structure(_get_tensors_for_gradient, gradient_components)
+
+
+def _replace_tensors_for_gradient(x, grad):
+  """Replaces the tensors in `x` that should be differentiated with `grad`.
+
+  Args:
+    x: A `Tensor` or `CompositeTensor`.
+    grad: A nested structure of `Tensor`, with the same structure as the value
+      returned by `_get_tensors_for_gradient(x)`.
+
+  Returns:
+    A `Tensor` or `CompositeTensor`.
+  """
+  if not isinstance(x, composite_tensor.CompositeTensor):
+    return grad
+
+  if not isinstance(x, CompositeTensorGradientProtocol):
+    raise ValueError(
+        f"Type {type(x).__name__} is not supported as a gradient source.")
+
+  composite_gradient = x.__composite_gradient__
+  x_components = composite_gradient.get_gradient_components(x)
+  if x_components is x:
+    grad_components = grad
+  else:
+    grad_components = nest.map_structure_up_to(x_components,
+                                               _replace_tensors_for_gradient,
+                                               x_components, grad)
+  if grad_components is None:
+    return None
+  return composite_gradient.replace_gradient_components(x, grad_components)
+
+
+def get_flat_tensors_for_gradients(xs):
+  """Returns a flat list of Tensors that should be differentiated for `xs`.
+
+  Args:
+    xs: A list of `Tensor`s or `CompositeTensor`s.
+
+  Returns:
+    A flat list of `Tensor`s constructed from `xs`, where `Tensor` values are
+    left as-is, and `CompositeTensor`s are replaced with
+    `_get_tensors_for_gradient(x)`.
+  """
+  return nest.flatten([_get_tensors_for_gradient(x) for x in xs])
+
+
+def replace_flat_tensors_for_gradients(xs, flat_grads):
+  """Replaces Tensors that should be differentiated in `xs` with `flat_grads`.
+
+  Args:
+    xs: A list of `Tensor`s or `CompositeTensor`s.
+    flat_grads: A list of `Tensor`.
+
+  Returns:
+    A list of `Tensor` or `CompositeTensor`.
+  """
+  xs_structure = [_get_tensors_for_gradient(x) for x in xs]
+  grads = nest.pack_sequence_as(xs_structure, flat_grads)
+  return [_replace_tensors_for_gradient(x, grad) for x, grad in zip(xs, grads)]

tensorflow/python/framework/composite_tensor_test.py

@@ -0,0 +1,429 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.composite_tensor."""
+
+import gc
+import sys
+import weakref
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+from tensorflow.python.util import nest
+
+
+class CTSpec(type_spec.TypeSpec):
+  """A generic CompositeTensor TypeSpec, used for constructing tests."""
+
+  def __init__(self, component_specs, metadata=None):
+    self.component_specs = component_specs
+    self.metadata = metadata
+
+  value_type = property(lambda self: CT)
+  _component_specs = property(lambda self: self.component_specs)
+
+  def _serialize(self):
+    return (self.component_specs, self.metadata)
+
+  def _to_components(self, value):
+    return value.components
+
+  def _from_components(self, tensor_list):
+    return CT(tensor_list, self.metadata)
+
+
+class CT(composite_tensor.CompositeTensor):
+  """A generic CompositeTensor, used for constructing tests."""
+  _type_spec_class = CTSpec
+
+  def __init__(self, components, metadata=None):
+    if isinstance(components, list):
+      components = tuple(components)
+    self.components = components
+    self.metadata = metadata
+
+  @property
+  def _type_spec(self):
+    component_specs = nest.map_structure(type_spec.type_spec_from_value,
+                                         self.components)
+    return self._type_spec_class(component_specs, self.metadata)
+
+  def __repr__(self):
+    return '%s(%r, %r)' % (type(self).__name__, self.components, self.metadata)
+
+  def __eq__(self, other):
+    return (type(self) is type(other) and
+            self.components == other.components and
+            self.metadata == other.metadata)
+
+
+# Another test CompositeTensor class.  `tf.nest` should treat different CT
+# classes without common supertypes as different structure types
+# (e.g. for assert_same_structure).
+class CTSpec2(CTSpec):
+  pass
+
+
+class CT2(CT):
+  _type_spec_class = CTSpec2
+
+
+# CompositeTensors with a common supertype are considered to be the same
+# structure by tf.nest (e.g. for assert_same_structure).
+class CT3(CT):
+  _type_spec_class = CTSpec
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class CompositeTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      {'structure': CT(0),
+       'expected': [0],
+       'paths': [('CT',)]},
+      {'structure': CT('a'),
+       'expected': ['a'],
+       'paths': [('CT',)]},
+      {'structure': CT(['a', 'b', 'c']),
+       'expected': ['a', 'b', 'c'],
+       'paths': [('CT', 0), ('CT', 1), ('CT', 2)]},
+      {'structure': CT({'x': 'a', 'y': 'b', 'z': 'c'}),
+       'expected': ['a', 'b', 'c'],
+       'paths': [('CT', 'x'), ('CT', 'y'), ('CT', 'z')]},
+      {'structure': [{'k1': CT('a')}, CT(['b', {'x': CT({'y': 'c'})}])],
+       'expected': ['a', 'b', 'c'],
+       'paths': [(0, 'k1', 'CT'), (1, 'CT', 0), (1, 'CT', 1, 'x', 'CT', 'y')]},
+      {'structure': CT(0),
+       'expand_composites': False,
+       'expected': [CT(0)],
+       'paths': [()]},
+      {'structure': [{'k1': CT('a')}, CT(['b', {'x': CT({'y': 'c'})}])],
+       'expand_composites': False,
+       'expected': [CT('a'), CT(['b', {'x': CT({'y': 'c'})}])],
+       'paths': [(0, 'k1'), (1,)]},
+  ])  # pyformat: disable
+  def testNestFlatten(self, structure, expected, paths, expand_composites=True):
+    result = nest.flatten(structure, expand_composites=expand_composites)
+    self.assertEqual(result, expected)
+
+    result_with_paths = nest.flatten_with_tuple_paths(
+        structure, expand_composites=expand_composites)
+    self.assertEqual(result_with_paths, list(zip(paths, expected)))
+
+    string_paths = ['/'.join(str(p) for p in path) for path in paths]  # pylint: disable=g-complex-comprehension
+    result_with_string_paths = nest.flatten_with_joined_string_paths(
+        structure, expand_composites=expand_composites)
+    self.assertEqual(result_with_string_paths,
+                     list(zip(string_paths, expected)))
+
+    flat_paths_result = list(
+        nest.yield_flat_paths(structure, expand_composites=expand_composites))
+    self.assertEqual(flat_paths_result, paths)
+
+  @parameterized.parameters([
+      {'s1': [1, 2, 3],
+       's2': [CT(['a', 'b']), 'c', 'd'],
+       'expand_composites': False,
+       'expected': [CT(['a', 'b']), 'c', 'd'],
+       'paths': [(0,), (1,), (2,)]},
+      {'s1': [CT([1, 2, 3])],
+       's2': [5],
+       'expand_composites': False,
+       'expected': [5],
+       'paths': [(0,)]},
+      {'s1': [[CT([9, 9, 9])], 999, {'y': CT([9, 9])}],
+       's2': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       'expand_composites': False,
+       'expected': [CT([1, 2, 3]), 100, CT([CT([4, 5]), 6])],
+       'paths': [(0, 0), (1,), (2, 'y')]},
+      {'s1': [[CT([9, 9, 9])], 999, {'y': CT([CT([9, 9]), 9])}],
+       's2': [[CT([1, 2, 3])], 100, {'y': CT([5, 6])}],
+       'expand_composites': False,
+       'expected': [CT([1, 2, 3]), 100, CT([5, 6])],
+       'paths': [(0, 0), (1,), (2, 'y')]},
+  ])  # pyformat: disable
+  def testNestFlattenUpTo(self, s1, s2, expected, paths,
+                          expand_composites=True):
+    result = nest.flatten_up_to(s1, s2, expand_composites=expand_composites)
+    self.assertEqual(expected, result)
+
+    result_with_paths = nest.flatten_with_tuple_paths_up_to(
+        s1, s2, expand_composites=expand_composites)
+    self.assertEqual(result_with_paths, list(zip(paths, expected)))
+
+  @parameterized.parameters([
+      {'structure': CT(0),
+       'sequence': [5],
+       'expected': CT(5)},
+      {'structure': CT(['a', 'b', 'c']),
+       'sequence': ['A', CT(['b']), {'x': 'y'}],
+       'expected': CT(['A', CT(['b']), {'x': 'y'}])},
+      {'structure': [{'k1': CT('a')}, CT(['b', {'x': CT({'y': 'c'})}])],
+       'sequence': ['A', 'B', 'C'],
+       'expected': [{'k1': CT('A')}, CT(['B', {'x': CT({'y': 'C'})}])]},
+      {'structure': [{'k1': CT('a')}, CT(['b', {'x': CT({'y': 'c'})}])],
+       'sequence': ['A', 'B'],
+       'expand_composites': False,
+       'expected': [{'k1': 'A'}, 'B']},
+      {'structure': CT(0, metadata='abc'),
+       'sequence': [5],
+       'expected': CT(5, metadata='abc')},
+  ])  # pyformat: disable
+  def testNestPackSequenceAs(self,
+                             structure,
+                             sequence,
+                             expected,
+                             expand_composites=True):
+    result = nest.pack_sequence_as(
+        structure, sequence, expand_composites=expand_composites)
+    self.assertEqual(result, expected)
+
+  @parameterized.parameters([
+      {'s1': CT('abc'), 's2': CT('xyz')},
+      {'s1': CT(['a', 'b', 'c']), 's2': CT(['d', 'e', 'f'])},
+      {'s1': [1, CT([10]), CT(200, metadata='xyz')],
+       's2': [8, CT([55]), CT(100, metadata='xyz')]},
+      {'s1': CT('abc'), 's2': CT3('xyz')},
+      {'s1': CT(['a', 'b', 'c']), 's2': CT3(['d', 'e', 'f'])},
+      {'s1': [1, CT([10]), CT(200, metadata='xyz')],
+       's2': [8, CT([55]), CT3(100, metadata='xyz')]},
+  ])  # pyformat: disable
+  def testNestAssertSameStructure(self, s1, s2, expand_composites=True):
+    nest.assert_same_structure(s1, s2, expand_composites=expand_composites)
+    nest.assert_shallow_structure(s1, s2, expand_composites=expand_composites)
+
+  @parameterized.parameters([
+      {'s1': CT(0), 's2': CT(['x'])},
+      {'s1': CT([1]), 's2': CT([1, 2])},
+      {'s1': CT({'x': 1}), 's2': CT({'y': 1})},
+      {'s1': CT(0), 's2': CT(0, metadata='xyz')},
+      {'s1': CT(0, metadata='xyz'), 's2': CT(0)},
+      {'s1': CT(0, metadata='xyz'), 's2': CT(0, metadata='abc')},
+      {'s1': CT(['a', 'b', 'c']), 's2': CT(['d', 'e'])},
+      {'s1': [1, CT(['a']), CT('b', metadata='xyz')],
+       's2': [8, CT([55, 66]), CT(100, metadata='abc')]},
+      {'s1': CT(0), 's2': CT2(0)},
+  ])  # pyformat: disable
+  def testNestAssertSameStructureCompositeMismatch(self,
+                                                   s1,
+                                                   s2,
+                                                   error=ValueError):
+    # s1 and s2 have the same structure if expand_composites=False; but
+    # different structures if expand_composites=True.
+    nest.assert_same_structure(s1, s2, expand_composites=False)
+    nest.assert_shallow_structure(s1, s2, expand_composites=False)
+    with self.assertRaises(error):  # pylint: disable=g-error-prone-assert-raises
+      nest.assert_same_structure(s1, s2, expand_composites=True)
+
+  @parameterized.parameters([
+      # Note: there are additional test cases in testNestAssertSameStructure.
+      {'s1': [1], 's2': [CT(1)]},
+      {'s1': [[CT([1, 2, 3])], 100, {'y': CT([5, 6])}],
+       's2': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       'expand_composites': False},
+      {'s1': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       's2': [[CT([1, 2, 3])], 100, {'y': CT([5, 6])}],
+       'expand_composites': False},
+  ])  # pyformat: disable
+  def testNestAssertShallowStructure(self, s1, s2, expand_composites=True):
+    nest.assert_shallow_structure(s1, s2, expand_composites=expand_composites)
+
+  @parameterized.parameters([
+      # Note: there are additional test cases in
+      # testNestAssertSameStructureCompositeMismatch.
+      {'s1': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       's2': [[CT([1, 2, 3])], 100, {'y': CT([5, 6])}]},
+      {'s1': CT([1, 2, 3]),
+       's2': [1, 2, 3],
+       'check_types': False},
+  ])  # pyformat: disable
+  def testNestAssertShallowStructureCompositeMismatch(self,
+                                                      s1,
+                                                      s2,
+                                                      check_types=True):
+    with self.assertRaises((TypeError, ValueError)):  # pylint: disable=g-error-prone-assert-raises
+      nest.assert_shallow_structure(
+          s1, s2, expand_composites=True, check_types=check_types)
+
+  @parameterized.parameters([
+      {'structure': CT(1, metadata=2),
+       'expected': CT(11, metadata=2)},
+      {'structure': CT({'x': 1, 'y': [2, 3]}, metadata=2),
+       'expected': CT({'x': 11, 'y': [12, 13]}, metadata=2)},
+      {'structure': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       'expected': [[CT([11, 12, 13])], 110, {'y': CT([CT([14, 15]), 16])}]},
+  ])  # pyformat: disable
+  def testNestMapStructure(self, structure, expected, expand_composites=True):
+    func = lambda x: x + 10
+    result = nest.map_structure(
+        func, structure, expand_composites=expand_composites)
+    self.assertEqual(result, expected)
+
+  @parameterized.parameters([
+      {'s1': [[CT([1, 2, 3])], 100, {'y': 4}],
+       's2': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       'expected': [[CT([11, 12, 13])], 110, {'y': CT([CT([4, 5]), 6])}]}
+  ])  # pyformat: disable
+  def testNestMapStructureUpTo(self, s1, s2, expected):
+    func = lambda x: x + 10 if isinstance(x, int) else x
+    result = nest.map_structure_up_to(s1, func, s2, expand_composites=True)
+    self.assertEqual(result, expected)
+
+  @parameterized.parameters([
+      {'structure': CT('a'),
+       'expected': CT('CT:a')},
+      {'structure': CT(['a', 'b']),
+       'expected': CT(['CT/0:a', 'CT/1:b'])},
+      {'structure': [[CT([1, 2, 3])], 100, {'y': CT([CT([4, 5]), 6])}],
+       'expected': [
+           [CT(['0/0/CT/0:1', '0/0/CT/1:2', '0/0/CT/2:3'])],
+           '1:100',
+           {'y': CT([CT(['2/y/CT/0/CT/0:4', '2/y/CT/0/CT/1:5']),
+                     '2/y/CT/1:6'])}]},
+  ])  # pyformat: disable
+  def testNestMapStructureWithPaths(self,
+                                    structure,
+                                    expected,
+                                    expand_composites=True):
+
+    def func1(path, x):
+      return '%s:%s' % (path, x)
+
+    result = nest.map_structure_with_paths(
+        func1, structure, expand_composites=expand_composites)
+    self.assertEqual(result, expected)
+
+    # Use the same test cases for map_structure_with_tuple_paths.
+    def func2(tuple_path, x):
+      return '%s:%s' % ('/'.join(str(v) for v in tuple_path), x)
+
+    result = nest.map_structure_with_tuple_paths(
+        func2, structure, expand_composites=expand_composites)
+    self.assertEqual(result, expected)
+
+  @parameterized.parameters([
+      {'s1': [[CT([1, 2, 3])], 100, {'y': [4, 5]}],
+       's2': [[CT([1, 2, 3])], 100, {'y': [CT([4, 5]), 6]}],
+       'expected': [
+           [CT(['0/0/CT/0:1', '0/0/CT/1:2', '0/0/CT/2:3'])],
+           ('1:100'),
+           {'y': ['2/y/0:CT((4, 5), None)', '2/y/1:6']}]},
+  ])  # pyformat: disable
+  def testNestMapStructureWithTuplePathsUpTo(self, s1, s2, expected):
+
+    def func(tuple_path, x):
+      return '%s:%s' % ('/'.join(str(v) for v in tuple_path), x)
+
+    result = nest.map_structure_with_tuple_paths_up_to(
+        s1, func, s2, expand_composites=True)
+    self.assertEqual(result, expected)
+
+  def testNestGetTraverseShallowStructure(self):
+    func = lambda t: not (isinstance(t, CT) and t.metadata == 'B')
+    structure = [CT([1, 2], metadata='A'), CT([CT(3)], metadata='B')]
+
+    result = nest.get_traverse_shallow_structure(
+        func, structure, expand_composites=True)
+    expected = [CT([True, True], metadata='A'), False]
+    self.assertEqual(result, expected)
+
+  def testMemoryIsFreed(self):
+    # Note: we use `np.array` values for CT and `set` values for
+    # metadata because we need to construct weakrefs to them.  Other builtin
+    # types, such as `list` and `tuple`, do not support weakrefs.
+    ct1 = CT(np.array([1, 2]), set(['no', 'leaks']))
+    ct2 = CT(np.array([3, 4]), set(['no', 'leaks']))
+    ct3 = CT(np.array([5, 6]), set(['other', 'metadata']))
+
+    # Note: map_structure exercises flatten, pack_sequence_as, and
+    # assert_same_structure.
+    func = lambda x, y: x + y
+    ct4 = nest.map_structure(func, ct1, ct2, expand_composites=True)
+
+    # Check that the exception-raising path in assert_same_structure
+    # doesn't leak any objects.
+    with self.assertRaises(ValueError):
+      nest.map_structure(func, ct2, ct3, expand_composites=True)
+    if hasattr(sys, 'exc_clear'):
+      sys.exc_clear()  # Remove any references in exception stack traces.
+
+    refs = []
+    for ct in [ct1, ct2, ct3, ct4]:
+      refs.append(weakref.ref(ct))
+      refs.append(weakref.ref(ct.components))
+      refs.append(weakref.ref(ct.metadata))
+    del ct  # pylint: disable=undefined-loop-variable
+
+    for ref in refs:
+      self.assertIsNotNone(ref())
+
+    del ct1, ct2, ct3, ct4
+    gc.collect()
+    for ref in refs:
+      self.assertIsNone(ref())
+
+  # pylint: disable=g-long-lambda
+  @parameterized.named_parameters([
+      ('IndexedSlicesNoDenseShape', lambda: indexed_slices.IndexedSlices(
+          constant_op.constant([1, 2, 3]), constant_op.constant([2, 8, 4]))),
+      ('IndexedSlicesInt32DenseShape', lambda: indexed_slices.IndexedSlices(
+          constant_op.constant([1, 2, 3]), constant_op.constant([2, 8, 4]),
+          constant_op.constant([10], dtypes.int32))),
+      ('IndexedSlicesInt64DenseShape', lambda: indexed_slices.IndexedSlices(
+          constant_op.constant([[1, 2], [3, 4]]), constant_op.constant([2, 8]),
+          constant_op.constant([10, 2], dtypes.int64))),
+      ('RaggedTensorRaggedRank1',
+       lambda: ragged_factory_ops.constant([[1, 2], [3]])),
+      ('RaggedTensorRaggedRank2',
+       lambda: ragged_factory_ops.constant([[[1, 2], [3]], [[6, 7, 8]]])),
+      ('SparseTensor',
+       lambda: sparse_tensor.SparseTensor([[3], [7]], ['a', 'b'], [10])),
+      ('Nested structure', lambda: {
+          'a':
+              indexed_slices.IndexedSlices(
+                  constant_op.constant([1, 2, 3]),
+                  constant_op.constant([2, 8, 4])),
+          'b': [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              sparse_tensor.SparseTensor([[3], [7]], ['a', 'b'], [10])
+          ]
+      }),
+  ])
+  def testAssertSameStructureWithValueAndTypeSpec(self, value_func):
+    value = value_func()
+    spec = nest.map_structure(type_spec.type_spec_from_value, value,
+                              expand_composites=False)
+    nest.assert_same_structure(value, spec, expand_composites=True)
+
+  def testConvertVariablesToTensors(self):
+    ct = CT(1)
+    result = ct._convert_variables_to_tensors()
+    self.assertIs(result, ct)
+
+    result2 = composite_tensor.convert_variables_to_tensors(ct)
+    self.assertIs(result2, ct)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/framework/config_test.py

@@ -0,0 +1,912 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests that the system configuration methods work properly."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+def reset_eager(fn):
+
+  def wrapper(*args, **kwargs):
+    try:
+      return fn(*args, **kwargs)
+    finally:
+      # Reset the context.
+      context._reset_jit_compiler_flags()
+      context._reset_context()
+      ops.enable_eager_execution_internal()
+      assert context._context is not None
+
+  return wrapper
+
+
+@test_util.with_eager_op_as_function
+class ConfigTest(test.TestCase, parameterized.TestCase):
+
+  @test_util.disable_eager_op_as_function('b/204320409')
+  @test_util.run_gpu_only
+  @reset_eager
+  def testDevicePolicy(self):
+    self.assertEqual(context.DEVICE_PLACEMENT_SILENT,
+                     context.context().device_policy)
+
+    # If no op has been executed we should be able to set the device policy as
+    # well as any init-time configs.
+    config.set_intra_op_parallelism_threads(1)
+    config.set_device_policy('silent')
+    config.set_intra_op_parallelism_threads(2)
+
+    context.ensure_initialized()
+
+    def copy_tensor(dtype=dtypes.int32):
+      with ops.device('CPU:0'):
+        cpu_tensor = constant_op.constant(1, dtype=dtype)
+      gpu_tensor = cpu_tensor.gpu()
+      self.assertAllEqual(cpu_tensor + gpu_tensor, 2.0)
+
+    config.set_device_policy('silent')
+    self.assertEqual(config.get_device_policy(), 'silent')
+    self.assertEqual(context.DEVICE_PLACEMENT_SILENT,
+                     context.context().device_policy)
+    copy_tensor()
+
+    config.set_device_policy('silent_for_int32')
+    self.assertEqual(config.get_device_policy(), 'silent_for_int32')
+    self.assertEqual(context.DEVICE_PLACEMENT_SILENT_FOR_INT32,
+                     context.context().device_policy)
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Tensors on conflicting devices'):
+      copy_tensor(dtypes.float32)
+    copy_tensor()
+
+    config.set_device_policy('warn')
+    self.assertEqual(config.get_device_policy(), 'warn')
+    self.assertEqual(context.DEVICE_PLACEMENT_WARN,
+                     context.context().device_policy)
+    copy_tensor()
+
+    config.set_device_policy('explicit')
+    self.assertEqual(config.get_device_policy(), 'explicit')
+    self.assertEqual(context.DEVICE_PLACEMENT_EXPLICIT,
+                     context.context().device_policy)
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Tensors on conflicting devices'):
+      copy_tensor()
+
+    config.set_device_policy(None)
+    self.assertEqual(config.get_device_policy(), 'silent')
+
+  @reset_eager
+  def testExecutionMode(self):
+    self.assertTrue(config.get_synchronous_execution())
+    self.assertEqual(context.SYNC, context.context().execution_mode)
+
+    # If no op has been executed we should be able to set the execution mode as
+    # well as any init-time configs.
+    config.set_intra_op_parallelism_threads(1)
+    config.set_synchronous_execution(False)
+    config.set_intra_op_parallelism_threads(2)
+
+    config.set_synchronous_execution(True)
+    self.assertTrue(config.get_synchronous_execution())
+    self.assertEqual(context.SYNC, context.context().execution_mode)
+    config.set_synchronous_execution(False)
+    self.assertFalse(config.get_synchronous_execution())
+    self.assertEqual(context.ASYNC, context.context().execution_mode)
+
+  @reset_eager
+  def testIntraOpParallelismThreads(self):
+    config.set_intra_op_parallelism_threads(10)
+    self.assertEqual(config.get_intra_op_parallelism_threads(),
+                     context.context().intra_op_parallelism_threads)
+
+    context.ensure_initialized()
+
+    with self.assertRaises(RuntimeError):
+      config.set_intra_op_parallelism_threads(1)
+
+    config.set_intra_op_parallelism_threads(10)
+
+  @reset_eager
+  def testInterOpParallelismThreads(self):
+    config.set_inter_op_parallelism_threads(10)
+    self.assertEqual(config.get_inter_op_parallelism_threads(),
+                     context.context().inter_op_parallelism_threads)
+
+    context.ensure_initialized()
+
+    with self.assertRaises(RuntimeError):
+      config.set_inter_op_parallelism_threads(1)
+
+    config.set_inter_op_parallelism_threads(10)
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testSoftPlacement(self):
+    if context.executing_eagerly():
+      self.assertTrue(config.get_soft_device_placement())
+    else:
+      self.assertFalse(config.get_soft_device_placement())
+
+    def test_attr():
+      with ops.device('/device:GPU:0'):
+        return test_ops.test_attr(T=dtypes.float32, name='test_attr')
+
+    config.set_soft_device_placement(True)
+    self.assertEqual(config.get_soft_device_placement(), True)
+    self.assertEqual(config.get_soft_device_placement(),
+                     context.context().soft_device_placement)
+
+    # Since soft placement is enabled, the test_attr operation should fallback
+    # to CPU with pure eager execution as well as functions
+    test_attr()
+    def_function.function(test_attr)()
+
+    config.set_soft_device_placement(False)
+    self.assertEqual(config.get_soft_device_placement(), False)
+    self.assertEqual(config.get_soft_device_placement(),
+                     context.context().soft_device_placement)
+
+    # Since soft placement is disabled, the test_attr operation should fail on
+    # GPU with pure eager execution as well as functions
+    with self.assertRaises(errors.InvalidArgumentError):
+      test_attr()
+    with self.assertRaises(errors.InvalidArgumentError):
+      def_function.function(test_attr)()
+
+  @reset_eager
+  def testLogDevicePlacement(self):
+    self.assertFalse(context.get_log_device_placement())
+
+    context.set_log_device_placement(True)
+    self.assertEqual(context.get_log_device_placement(), True)
+    self.assertEqual(context.get_log_device_placement(),
+                     context.context().log_device_placement)
+
+    context.set_log_device_placement(False)
+    self.assertEqual(context.get_log_device_placement(), False)
+    self.assertEqual(context.get_log_device_placement(),
+                     context.context().log_device_placement)
+
+    context.ensure_initialized()
+
+    # Changing the device placement should not throw an exception
+    context.set_log_device_placement(True)
+
+  @reset_eager
+  def testEnableMlirBridge(self):
+    # Default value of enable_mlir_bridge is false.
+    self.assertFalse(context.context().config.experimental.enable_mlir_bridge)
+    self.assertEqual(
+        context.context().config.experimental.mlir_bridge_rollout,
+        config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_UNSPECIFIED)
+
+    # Tests enabling mlir bridge.
+    config.enable_mlir_bridge()
+    self.assertTrue(context.context().config.experimental.enable_mlir_bridge)
+    self.assertEqual(
+        context.context().config.experimental.mlir_bridge_rollout,
+        config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_ENABLED)
+
+    # Tests disabling mlir bridge.
+    config.disable_mlir_bridge()
+    self.assertFalse(context.context().config.experimental.enable_mlir_bridge)
+    self.assertEqual(
+        context.context().config.experimental.mlir_bridge_rollout,
+        config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_DISABLED)
+
+  @reset_eager
+  def testResetMlirFlags(self):
+    # Default value of enable_mlir_bridge is false.
+    self.assertFalse(context.context().config.experimental.enable_mlir_bridge)
+    self.assertEqual(
+        context.context().config.experimental.mlir_bridge_rollout,
+        config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_UNSPECIFIED)
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testJit(self):
+    self.assertEqual(config.get_optimizer_jit(), '')
+
+    # the following function should cause Op fusion to occur. However, there is
+    # unfortunately no straightforward way to ensure this. We will just have to
+    # settle for creating a test that can trigger JIT.
+    @def_function.function
+    def fun(a, b):
+      c = a * b
+      d = c + a
+      return d
+
+    a = constant_op.constant([2., 2.])
+    b = constant_op.constant([2., 2.])
+
+    self.evaluate(fun(a, b))
+
+    config.set_optimizer_jit('autoclustering')
+    self.assertEqual(config.get_optimizer_jit(), 'autoclustering')
+
+    self.evaluate(fun(a, b))
+
+    config.set_optimizer_jit('')
+    self.assertEqual(config.get_optimizer_jit(), '')
+
+    self.evaluate(fun(a, b))
+
+  @parameterized.named_parameters(
+      ('LayoutOptimizer', 'layout_optimizer'),
+      ('ConstantFolding', 'constant_folding'),
+      ('ShapeOptimization', 'shape_optimization'), ('Remapping', 'remapping'),
+      ('ArithmeticOptimization', 'arithmetic_optimization'),
+      ('DependencyOptimization', 'dependency_optimization'),
+      ('LoopOptimization', 'loop_optimization'),
+      ('FunctionOptimization', 'function_optimization'),
+      ('DebugStripper', 'debug_stripper'),
+      ('ScopedAllocatorOptimization', 'scoped_allocator_optimization'),
+      ('ImplementationSelector', 'implementation_selector'),
+      ('AutoMixedPrecision', 'auto_mixed_precision'))
+  @reset_eager
+  def testOptimizerToggleOption(self, field):
+    # TODO(b/128531235): Improve testing of option
+    options = config.get_optimizer_experimental_options()
+    self.assertIsNone(options.get(field))
+
+    config.set_optimizer_experimental_options({field: True})
+    options[field] = True
+    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
+    self.assertDictEqual(context.context().get_optimizer_experimental_options(),
+                         options)
+
+    config.set_optimizer_experimental_options({field: False})
+    options[field] = False
+    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
+    self.assertDictEqual(context.context().get_optimizer_experimental_options(),
+                         options)
+
+  @parameterized.named_parameters(
+      ('DisableModelPruning', 'disable_model_pruning'),
+      ('DisableMetaOptimizer', 'disable_meta_optimizer'))
+  @reset_eager
+  def testOptimizerBoolOption(self, field):
+    # TODO(b/128531235): Improve testing of option
+    options = config.get_optimizer_experimental_options()
+    self.assertFalse(options.get(field))
+
+    config.set_optimizer_experimental_options({field: True})
+    options[field] = True
+    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
+    self.assertDictEqual(context.context().get_optimizer_experimental_options(),
+                         options)
+
+    config.set_optimizer_experimental_options({field: False})
+    options[field] = False
+    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
+    self.assertDictEqual(context.context().get_optimizer_experimental_options(),
+                         options)
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testOptimizerToggleOptionPinToHost(self):
+    options = config.get_optimizer_experimental_options()
+    self.assertIsNone(options.get('pin_to_host_optimization'))
+
+    @def_function.function
+    def fun():
+      op = test_ops.device_placement_op()
+      return op
+
+    # Force optimizer to run for all graphs
+    config.set_optimizer_experimental_options({'min_graph_nodes': -1})
+    options['min_graph_nodes'] = -1
+
+    # Since pin to host is disabled, the operation should go on GPU
+    gpu = self.evaluate(fun())
+    self.assertIn(compat.as_bytes('GPU'), gpu)
+
+    config.set_optimizer_experimental_options(
+        {'pin_to_host_optimization': True})
+    options['pin_to_host_optimization'] = True
+    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
+    self.assertDictEqual(context.context().get_optimizer_experimental_options(),
+                         options)
+
+    # Since pin to host is enabled, the operation should go on CPU
+    cpu = self.evaluate(fun())
+    self.assertIn(compat.as_bytes('CPU'), cpu)
+
+    config.set_optimizer_experimental_options(
+        {'pin_to_host_optimization': False})
+    options['pin_to_host_optimization'] = False
+    self.assertDictEqual(config.get_optimizer_experimental_options(), options)
+    self.assertDictEqual(context.context().get_optimizer_experimental_options(),
+                         options)
+
+    # Since pin to host is disabled again, the operation should go on GPU
+    gpu2 = self.evaluate(fun())
+    self.assertIn(compat.as_bytes('GPU'), gpu2)
+
+
+class DeviceTest(test.TestCase):
+
+  @reset_eager
+  def testPhysicalDevices(self):
+    cpus = config.list_physical_devices('CPU')
+    self.assertGreater(len(cpus), 0)
+    if test_util.is_gpu_available():
+      gpus = config.list_physical_devices('GPU')
+      self.assertGreater(len(gpus), 0)
+
+  @reset_eager
+  def testCpuMultiple(self):
+    cpus = config.list_physical_devices('CPU')
+    self.assertEqual(len(cpus), 1)
+
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+
+    context.ensure_initialized()
+
+    vcpus = config.list_logical_devices('CPU')
+    self.assertEqual(len(vcpus), 2)
+
+    with ops.device('/device:CPU:0'):
+      a = constant_op.constant(1.0)
+      self.evaluate(a)
+    with ops.device('/device:CPU:1'):
+      b = constant_op.constant(1.0)
+      self.evaluate(b)
+    with ops.device('/device:CPU:2'):
+      c = constant_op.constant(1.0)
+      self.evaluate(c)
+    if test_util.is_gpu_available():
+      self.assertIn('GPU:0', c.device)
+    else:
+      self.assertIn('CPU:0', c.device)
+
+    # Ensure we can place ops on each of the device names
+    for vcpu in vcpus:
+      with ops.device(vcpu.name):
+        d = constant_op.constant(1.0)
+        self.evaluate(d)
+
+    # Modifying the CPU configuration is not supported
+    with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
+      config.set_logical_device_configuration(cpus[0], [
+          context.LogicalDeviceConfiguration(),
+          context.LogicalDeviceConfiguration(),
+          context.LogicalDeviceConfiguration()
+      ])
+
+    # Setting the same CPU configuration is fine
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testGpuNone(self):
+    config.set_soft_device_placement(False)
+    gpus = config.list_physical_devices('GPU')
+    self.assertGreater(len(gpus), 0)
+
+    cpus = config.list_physical_devices('CPU')
+    self.assertEqual(len(cpus), 1)
+
+    self.assertEqual(len(config.get_visible_devices('CPU')), 1)
+    self.assertGreater(len(config.get_visible_devices('GPU')), 0)
+
+    self.assertEqual(len(config.get_visible_devices('XLA_GPU')), 0)
+
+    config.set_visible_devices(cpus[0])
+    self.assertEqual(len(config.get_visible_devices('CPU')), 1)
+    self.assertEqual(len(config.get_visible_devices('GPU')), 0)
+    self.assertEqual(len(config.list_logical_devices('XLA_GPU')), 0)
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy'):
+      with ops.device('/device:GPU:0'):
+        a = array_ops.identity(1.0)
+        self.evaluate(a)
+
+    # Modifying the visible devices is not supported
+    with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
+      config.set_visible_devices(gpus)
+
+    # Setting the same visible devices is fine
+    config.set_visible_devices(cpus[0])
+
+  @reset_eager
+  def testGpuMultiple(self):
+    config.set_soft_device_placement(False)
+    gpus = config.list_physical_devices('GPU')
+
+    if len(gpus) < 2:
+      self.skipTest('Need at least 2 GPUs')
+
+    context.ensure_initialized()
+
+    for i in range(0, len(gpus)):
+      with ops.device('/device:GPU:' + str(i)):
+        a = constant_op.constant(1.0)
+        self.evaluate(a)
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      with ops.device('/device:GPU:' + str(len(gpus))):
+        a = constant_op.constant(1.0)
+        self.evaluate(a)
+
+  @reset_eager
+  def testDeviceDetails(self):
+    (cpu,) = config.list_physical_devices('CPU')
+    details = config.get_device_details(cpu)
+    self.assertEqual(details, {})
+
+    if not test_util.is_gpu_available():
+      return
+
+    gpus = config.list_physical_devices('GPU')
+    details = config.get_device_details(gpus[0])
+    self.assertIsInstance(details['device_name'], str)
+    self.assertNotEmpty(details['device_name'])
+    if test.is_built_with_rocm():
+      # AMD GPUs do not have a compute capability
+      self.assertNotIn('compute_capability', details)
+    else:
+      cc = details['compute_capability']
+      self.assertIsInstance(cc, tuple)
+      major, minor = cc
+      self.assertGreater(major, 0)
+      self.assertGreaterEqual(minor, 0)
+
+    # Test GPU returned from get_visible_devices
+    if len(gpus) > 2:
+      config.set_visible_devices(gpus[1], 'GPU')
+      (visible_gpu,) = config.get_visible_devices('GPU')
+      details = config.get_device_details(visible_gpu)
+      self.assertIsInstance(details['device_name'], str)
+
+  @reset_eager
+  def testDeviceDetailsErrors(self):
+    logical_devices = config.list_logical_devices()
+    with self.assertRaisesRegex(ValueError,
+                                'must be a tf.config.PhysicalDevice'):
+      config.get_device_details(logical_devices[0])
+
+    phys_dev = context.PhysicalDevice('/physical_device:CPU:100', 'CPU')
+    with self.assertRaisesRegex(
+        ValueError, 'The PhysicalDevice must be one obtained from '
+        'calling `tf.config.list_physical_devices`'):
+      config.get_device_details(phys_dev)
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testVirtualGpu(self):
+    config.set_soft_device_placement(False)
+    gpus = config.list_physical_devices('GPU')
+    self.assertNotEqual(len(gpus), 0)
+
+    self.assertIsNone(config.get_logical_device_configuration(gpus[-1]))
+    config.set_logical_device_configuration(gpus[-1], [
+        context.LogicalDeviceConfiguration(memory_limit=10),
+        context.LogicalDeviceConfiguration(memory_limit=10)
+    ])
+    self.assertEqual(len(config.get_logical_device_configuration(gpus[-1])), 2)
+
+    logical_gpus = config.list_logical_devices('GPU')
+    self.assertTrue(len(logical_gpus), len(gpus) + 1)
+    for i in range(0, len(logical_gpus)):
+      with ops.device('/device:GPU:' + str(i)):
+        a = array_ops.identity(1.0)
+        self.evaluate(a)
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy'):
+      with ops.device('/device:GPU:' + str(len(logical_gpus))):
+        a = array_ops.identity(1.0)
+        self.evaluate(a)
+
+    # Modifying the GPU configuration is not supported
+    with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
+      config.set_logical_device_configuration(gpus[-1], [
+          context.LogicalDeviceConfiguration(memory_limit=20),
+          context.LogicalDeviceConfiguration(memory_limit=20)
+      ])
+
+    with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
+      config.set_logical_device_configuration(gpus[-1], [
+          context.LogicalDeviceConfiguration(memory_limit=10),
+          context.LogicalDeviceConfiguration(memory_limit=10),
+          context.LogicalDeviceConfiguration(memory_limit=10)
+      ])
+
+    # Setting the same GPU configuration is fine
+    config.set_logical_device_configuration(gpus[-1], [
+        context.LogicalDeviceConfiguration(memory_limit=10),
+        context.LogicalDeviceConfiguration(memory_limit=10)
+    ])
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testGpuGrowth(self):
+    gpus = config.list_physical_devices('GPU')
+    self.assertNotEqual(len(gpus), 0)
+
+    self.assertIsNone(config.get_memory_growth(gpus[-1]))
+    for gpu in gpus:
+      config.set_memory_growth(gpu, True)
+
+    c = context.context().config
+    self.assertTrue(c.gpu_options.allow_growth)
+
+    logical_gpus = config.list_logical_devices('GPU')
+    self.assertTrue(len(logical_gpus), len(gpus))
+
+    # Modifying the GPU configuration is not supported
+    with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
+      for gpu in gpus:
+        config.set_memory_growth(gpu, False)
+
+    # Setting the same GPU configuration is fine
+    for gpu in gpus:
+      config.set_memory_growth(gpu, True)
+
+  @test_util.run_gpu_or_tpu
+  @reset_eager
+  def testGetMemoryInfoBasic(self, device_type):
+    with ops.device(f'{device_type}:0'):
+      device = array_ops.zeros([]).backing_device
+    info = config.get_memory_info(device)
+    self.assertGreater(info['current'], 0)
+    self.assertGreater(info['peak'], 0)
+    self.assertEqual(info.keys(), {'current', 'peak'})
+    self.assertEqual(config.get_memory_usage(device), info['current'])
+
+  @test_util.run_gpu_or_tpu
+  @reset_eager
+  def testGetMemoryUsageSubstring(self, device_type):
+    info = config.get_memory_info(f'{device_type}:0')
+    self.assertGreater(info['current'], 0)
+
+  @reset_eager
+  def testGetMemoryInfoCPU(self):
+    if test_util.IsMklEnabled():
+      # TODO(gzmkl) work with Google team to address design issue in allocator.h
+      self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
+
+    with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
+      config.get_memory_info('CPU:0')
+    with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
+      config.get_memory_usage('CPU:0')
+
+  @reset_eager
+  def testGetMemoryInfoUnknownDevice(self):
+    with self.assertRaisesRegex(ValueError, 'No matching devices found'):
+      config.get_memory_info('unknown_device:0')
+    with self.assertRaisesRegex(ValueError, 'No matching devices found'):
+      config.get_memory_usage('unknown_device:0')
+
+  @reset_eager
+  def testGetMemoryInfoInvalidDeviceString(self):
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU:')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU:CPU')
+
+  @test_util.run_gpu_or_tpu
+  @reset_eager
+  def testPeakMemoryUsage(self, device_type):
+    device = f'{device_type}:0'
+    with ops.device(device):
+      x1 = array_ops.zeros((1000, 1000))
+    peak1 = config.get_memory_info(device)['peak']
+    self.assertGreaterEqual(peak1, 4 * 1000 * 1000)
+    with ops.device(device):
+      x2 = array_ops.ones((1000, 1000))
+    peak2 = config.get_memory_info(device)['peak']
+    self.assertGreaterEqual(peak2, peak1 + 4 * 1000 * 1000)
+    del x1, x2  # With CPython, causes tensor memory to be immediately freed
+    peak3 = config.get_memory_info(device)['peak']
+    self.assertGreaterEqual(peak3, peak2)
+    self.assertGreaterEqual(peak3, config.get_memory_info(device)['current'])
+
+  @test_util.run_gpu_or_tpu
+  @reset_eager
+  def testResetMemoryStats(self, device_type):
+    device = f'{device_type}:0'
+    with ops.device(device):
+      x = array_ops.zeros((1000, 1000), dtype=dtypes.float32)
+    config.reset_memory_stats(device)
+    info1 = config.get_memory_info(device)
+    self.assertGreaterEqual(info1['peak'], 4 * 1000 * 1000)
+    self.assertGreaterEqual(info1['peak'], info1['current'])
+    self.assertGreater(info1['current'], 0)
+
+    del x  # With CPython, causes tensor memory to be immediately freed
+    config.reset_memory_stats(device)
+    info2 = config.get_memory_info(device)
+    self.assertLess(info2['peak'], info1['peak'])
+
+  @reset_eager
+  def testResetMemoryStatsCPU(self):
+    if test_util.IsMklEnabled():
+      # TODO(gzmkl) work with Google team to address design issue in allocator.h
+      self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
+
+    with self.assertRaisesRegex(ValueError, 'Cannot reset memory stats'):
+      config.reset_memory_stats('CPU:0')
+
+  @reset_eager
+  def testResetMemoryStatsUnknownDevice(self):
+    with self.assertRaisesRegex(ValueError, 'No matching devices found'):
+      config.reset_memory_stats('unknown_device:0')
+
+  @reset_eager
+  def testResetMemoryStatsInvalidDeviceString(self):
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().reset_memory_stats('GPU')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().reset_memory_stats('GPU:')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().reset_memory_stats('GPU:CPU')
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testGpuInvalidConfig(self):
+    gpus = config.list_physical_devices('GPU')
+    self.assertNotEqual(len(gpus), 0)
+
+    if len(gpus) > 1:
+      # Assert if other GPUs were not configured
+      config.set_memory_growth(gpus[0], True)
+      with self.assertRaisesRegex(ValueError, 'cannot differ'):
+        c = context.context().config
+
+      # If we limit visibility to GPU 0, growth is fine
+      config.set_visible_devices(gpus[0], 'GPU')
+      c = context.context().config
+      self.assertTrue(c.gpu_options.allow_growth)
+
+      # Default setting for second GPU is False and works if we set visibility
+      config.set_visible_devices(gpus[1], 'GPU')
+      c = context.context().config
+      self.assertFalse(c.gpu_options.allow_growth)
+
+      # Growth now fails because all the GPUs are visible and not the same
+      config.set_visible_devices(gpus, 'GPU')
+      with self.assertRaisesRegex(ValueError, 'cannot differ'):
+        c = context.context().config
+
+    for gpu in gpus:
+      config.set_memory_growth(gpu, True)
+
+    c = context.context().config
+    self.assertTrue(c.gpu_options.allow_growth)
+
+    with self.assertRaisesRegex(ValueError, 'memory limit'):
+      config.set_logical_device_configuration(gpus[-1], [
+          context.LogicalDeviceConfiguration(),
+          context.LogicalDeviceConfiguration()
+      ])
+
+    self.assertIsNone(config.get_logical_device_configuration(gpus[-1]))
+    config.set_logical_device_configuration(gpus[-1], [
+        context.LogicalDeviceConfiguration(memory_limit=10),
+        context.LogicalDeviceConfiguration(memory_limit=10)
+    ])
+
+    c = context.context().config
+    self.assertFalse(c.gpu_options.allow_growth)
+
+    with self.assertRaisesRegex(ValueError, 'virtual devices'):
+      config.set_memory_growth(gpus[-1], False)
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testRemote(self):
+    gpus = config.list_logical_devices('GPU')
+    self.assertNotEqual(len(gpus), 0)
+
+    context.ensure_initialized()
+
+    gpus = config.list_logical_devices('GPU')
+    self.assertNotEqual(len(gpus), 0)
+    for gpu in gpus:
+      self.assertIsNotNone(gpu.name)
+
+    context.ensure_initialized()
+
+    job_name = 'test'
+    cluster_def = cluster_pb2.ClusterDef()
+    job_def = cluster_def.job.add()
+    job_def.name = job_name
+    job_def.tasks[0] = 'localhost:0'
+
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_def, job_name=job_name, task_index=0, protocol='grpc')
+
+    context.set_server_def(server_def)
+
+    gpus = config.list_logical_devices('GPU')
+    for gpu in gpus:
+      self.assertIsNotNone(gpu.name)
+
+  @reset_eager
+  def testV1CompatibilityDummyInvisibleDeviceList(self):
+    gpus = config.list_physical_devices('GPU')
+    if gpus:
+      self.skipTest('Test requires no GPUs')
+
+    # Ensure GPU options left untouched on CPU only environments
+    context.context()._physical_devices = None
+    context.context()._config = config_pb2.ConfigProto(
+        gpu_options=config_pb2.GPUOptions(visible_device_list='0'))
+    new_config = context.context().config
+    self.assertEqual(new_config.gpu_options.visible_device_list, '0')
+
+  @test_util.run_gpu_only
+  @reset_eager
+  def testV1Compatibility(self):
+    # Ensure we set 1 CPU by default
+    context.context()._config = config_pb2.ConfigProto()
+    new_config = context.context().config
+    self.assertEqual(new_config.device_count['CPU'], 1)
+    context.context()._physical_devices = None
+
+    # Ensure CPU is split
+    context.context()._config = config_pb2.ConfigProto(device_count={'CPU': 2})
+    new_config = context.context().config
+    self.assertEqual(new_config.device_count['CPU'], 2)
+    context.context()._physical_devices = None
+
+    # Handle empty visible device list
+    context.context()._config = config_pb2.ConfigProto(
+        gpu_options=config_pb2.GPUOptions(visible_device_list=''))
+    gpus = config.list_physical_devices('GPU')
+    gpu_count = len(gpus)
+    new_config = context.context().config
+    self.assertEqual(new_config.gpu_options.visible_device_list,
+                     ','.join(str(i) for i in range(len(gpus))))
+    context.context()._physical_devices = None
+
+    # Handle invalid visible device list
+    context.context()._config = config_pb2.ConfigProto(
+        gpu_options=config_pb2.GPUOptions(visible_device_list=str(gpu_count)))
+    with self.assertRaisesRegex(ValueError, 'Invalid visible device index'):
+      gpus = config.list_physical_devices('GPU')
+      new_config = context.context().config
+    context.context()._physical_devices = None
+
+    # Handle single visible device list
+    context.context()._config = config_pb2.ConfigProto(
+        gpu_options=config_pb2.GPUOptions(
+            visible_device_list=str(gpu_count - 1)))
+    gpus = config.list_physical_devices('GPU')
+    new_config = context.context().config
+    self.assertEqual(new_config.gpu_options.visible_device_list,
+                     str(gpu_count - 1))
+    context.context()._physical_devices = None
+
+  def testConfigureCollectiveOps(self):
+    context.context().configure_collective_ops(
+        collective_leader='/job:worker/replica:0/task:0',
+        scoped_allocator_enabled_ops=('CollectiveReduce',),
+        use_nccl_communication=False,
+        device_filters=['/job:worker/task:1'])
+    new_config = context.context().config
+
+    # Verify group leader
+    self.assertEqual('/job:worker/replica:0/task:0',
+                     new_config.experimental.collective_group_leader)
+
+    # Verify device filters.
+    self.assertEqual(['/job:worker/task:1'], new_config.device_filters)
+
+    # Verify rewrite options.
+    new_rewrite_options = new_config.graph_options.rewrite_options
+    self.assertEqual(rewriter_config_pb2.RewriterConfig.ON,
+                     new_rewrite_options.scoped_allocator_optimization)
+    self.assertEqual(['CollectiveReduce'],
+                     new_rewrite_options.scoped_allocator_opts.enable_op)
+
+  def testDeterminism(self):
+    # This does not test any ops are deterministic, because that is tested by
+    # many kernel tests.
+    try:
+      config.disable_op_determinism()
+      self.assertFalse(config.is_op_determinism_enabled())
+      config.enable_op_determinism()
+      self.assertTrue(config.is_op_determinism_enabled())
+    finally:
+      config.disable_op_determinism()
+
+
+class TensorFloat32Test(test.TestCase):
+
+  def tearDown(self):
+    super(TensorFloat32Test, self).tearDown()
+    config.enable_tensor_float_32_execution(True)
+
+  def test_tensor_float_32_global_variable(self):
+    self.assertTrue(config.tensor_float_32_execution_enabled())
+    self.assertTrue(test_ops.is_tensor_float32_enabled())
+    config.enable_tensor_float_32_execution(False)
+    self.assertFalse(config.tensor_float_32_execution_enabled())
+    self.assertFalse(test_ops.is_tensor_float32_enabled())
+    config.enable_tensor_float_32_execution(True)
+    self.assertTrue(config.tensor_float_32_execution_enabled())
+    self.assertTrue(test_ops.is_tensor_float32_enabled())
+
+  def _skip_if_tensor_float_32_unsupported(self):
+    if not test_util.is_gpu_available(
+        cuda_only=True, min_cuda_compute_capability=(8, 0)):
+      self.skipTest('TensorFloat-32 requires an NVIDIA GPU with compute '
+                    'capability of at least 8.0')
+
+  # Size of each dimension of matrices to test. cuBLAS does not use TF32 for
+  # small matrices, so we must choose a large enough size to cause TF32 to be
+  # used.
+  DIM = 2 ** 10
+
+  def test_tensor_float_32_enabled(self):
+    self._skip_if_tensor_float_32_unsupported()
+    self.assertTrue(config.tensor_float_32_execution_enabled())
+
+    x = array_ops.fill((self.DIM, self.DIM), 1 + 2**-12)
+    y = array_ops.ones((self.DIM, self.DIM))
+    out = math_ops.matmul(x, y)
+    # In TensorFloat-32, each element of x is rounded to 1, so each output
+    # element should be self.DIM.
+    expected = array_ops.fill((self.DIM, self.DIM), float(self.DIM))
+    self.assertAllEqual(out, expected)
+
+  def test_tensor_float_32_disabled(self):
+    self._skip_if_tensor_float_32_unsupported()
+    self.assertTrue(config.tensor_float_32_execution_enabled())
+    config.enable_tensor_float_32_execution(False)
+    self.assertFalse(config.tensor_float_32_execution_enabled())
+
+    x = array_ops.fill((self.DIM, self.DIM), 1 + 2**-12)
+    y = array_ops.ones((self.DIM, self.DIM))
+    out = math_ops.matmul(x, y)
+    expected = array_ops.fill((self.DIM, self.DIM), self.DIM * (1 + 2**-12))
+    self.assertAllClose(out, expected, rtol=2**-13, atol=0)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/framework/config_vgpu_test.py

@@ -0,0 +1,136 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests that the system configuration methods work properly."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+def reset_eager(fn):
+
+  def wrapper(*args, **kwargs):
+    try:
+      return fn(*args, **kwargs)
+    finally:
+      # Reset the context.
+      context._reset_jit_compiler_flags()
+      context._reset_context()
+      ops.enable_eager_execution_internal()
+      assert context._context is not None
+
+  return wrapper
+
+
+class DeviceTest(test.TestCase):
+
+  # Due to a limitation in TensorFlow virtual GPU initialization, tests in
+  # this file must run with a new process per test. This is achieved by
+  # setting the number of shards to a large number.
+  already_run = False
+
+  def setUp(self):
+    super().setUp()
+    if self.already_run:
+      raise RuntimeError(
+          'Each test in this test suite must run in a separate process. '
+          'Increase number of shards used to run this test.')
+    self.already_run = True
+
+  @reset_eager
+  def testVGpu1P2V(self):
+    gpus = config.list_physical_devices('GPU')
+    if len(gpus) != 1:
+      self.skipTest('Need 1 GPUs')
+
+    config.set_logical_device_configuration(gpus[0], [
+        context.LogicalDeviceConfiguration(
+            memory_limit=100, experimental_device_ordinal=0),
+        context.LogicalDeviceConfiguration(
+            memory_limit=100, experimental_device_ordinal=1)
+    ])
+
+    context.ensure_initialized()
+
+    vcpus = config.list_logical_devices('GPU')
+    self.assertEqual(len(vcpus), 2)
+
+  @reset_eager
+  def testVGpu2P2V1Default(self):
+    gpus = config.list_physical_devices('GPU')
+    if len(gpus) != 2:
+      self.skipTest('Need 2 GPUs')
+
+    config.set_logical_device_configuration(gpus[0], [
+        context.LogicalDeviceConfiguration(memory_limit=100),
+        context.LogicalDeviceConfiguration(memory_limit=100)
+    ])
+
+    context.ensure_initialized()
+
+    vcpus = config.list_logical_devices('GPU')
+    self.assertEqual(len(vcpus), 3)
+
+  @reset_eager
+  def testGpu2P2V2V(self):
+    gpus = config.list_physical_devices('GPU')
+    if len(gpus) != 2:
+      self.skipTest('Need 2 GPUs')
+
+    config.set_logical_device_configuration(gpus[0], [
+        context.LogicalDeviceConfiguration(
+            memory_limit=100, experimental_device_ordinal=0),
+        context.LogicalDeviceConfiguration(
+            memory_limit=100, experimental_device_ordinal=1)
+    ])
+
+    config.set_logical_device_configuration(gpus[1], [
+        context.LogicalDeviceConfiguration(
+            memory_limit=100, experimental_device_ordinal=0),
+        context.LogicalDeviceConfiguration(
+            memory_limit=100, experimental_device_ordinal=1)
+    ])
+
+    context.ensure_initialized()
+
+    vcpus = config.list_logical_devices('GPU')
+    self.assertEqual(len(vcpus), 4)
+
+  @reset_eager
+  def testGpu2P2D2D(self):
+    gpus = config.list_physical_devices('GPU')
+    if len(gpus) != 2:
+      self.skipTest('Need 2 GPUs')
+
+    config.set_logical_device_configuration(gpus[0], [
+        context.LogicalDeviceConfiguration(memory_limit=100),
+        context.LogicalDeviceConfiguration(memory_limit=100)
+    ])
+
+    config.set_logical_device_configuration(gpus[1], [
+        context.LogicalDeviceConfiguration(memory_limit=100),
+        context.LogicalDeviceConfiguration(memory_limit=100)
+    ])
+
+    context.ensure_initialized()
+
+    vcpus = config.list_logical_devices('GPU')
+    self.assertEqual(len(vcpus), 4)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/framework/constant_op.py

@@ -0,0 +1,455 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Operations that generate constants.
+
+See the [constants guide](https://tensorflow.org/api_guides/python/constant_op).
+"""
+
+# Must be separate from array_ops to avoid a cyclic dependency.
+
+import numpy as np
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import execute
+# Import constant_tensor_conversion.py to register tensor conversion functions
+# for builtins. These functions were previously in this file, but were
+# refactored out so they can be registered at TF import time without importing
+# all of constant_op.py.
+from tensorflow.python.framework import constant_tensor_conversion  # pylint: disable=unused-import
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_conversion_registry
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.profiler import trace
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _eager_reshape(tensor, shape, ctx):
+  """Eager-only version of Reshape op; requires tensor is an eager Tensor."""
+  attr_t = tensor._datatype_enum()  # pylint: disable=protected-access
+  attr_tshape, (shape,) = execute.args_to_matching_eager(
+      [shape], ctx, [dtypes.int32, dtypes.int64], dtypes.int32)
+  inputs_flat = [tensor, shape]
+  attrs = ("T", attr_t, "Tshape", attr_tshape)
+  [result] = execute.execute(
+      b"Reshape", 1, inputs=inputs_flat, attrs=attrs, ctx=ctx)
+  return result
+
+
+def _eager_fill(dims, value, ctx):
+  """Eager-only version of Fill op; requires value is an eager Tensor."""
+  attr_t = value.dtype.as_datatype_enum
+  dims = convert_to_eager_tensor(dims, ctx, dtypes.int32)
+  inputs_flat = [dims, value]
+  attrs = ("T", attr_t, "index_type", types_pb2.DT_INT32)
+  [result] = execute.execute(
+      b"Fill", 1, inputs=inputs_flat, attrs=attrs, ctx=ctx)
+  return result
+
+
+def _eager_identity(tensor, ctx):
+  """Eager-only version of Identity op; requires tensor is an eager Tensor."""
+  attrs = ("T", tensor.dtype.as_datatype_enum)
+  [result] = execute.execute(
+      b"Identity", 1, inputs=[tensor], attrs=attrs, ctx=ctx)
+  return result
+
+
+def convert_to_eager_tensor(value, ctx, dtype=None) -> ops._EagerTensorBase:
+  """Converts the given `value` to an `EagerTensor`.
+
+  Note that this function could return cached copies of created constants for
+  performance reasons.
+
+  Args:
+    value: value to convert to EagerTensor.
+    ctx: value of context.context().
+    dtype: optional desired dtype of the converted EagerTensor.
+
+  Returns:
+    EagerTensor created from value.
+
+  Raises:
+    TypeError: if `dtype` is not compatible with the type of t.
+  """
+  if isinstance(value, np.ndarray):
+    # Make a copy explicitly because the EagerTensor might share the underlying
+    # memory with the input array. Without this copy, users will be able to
+    # modify the EagerTensor after its creation by changing the input array.
+    value = value.copy()
+  if isinstance(value, ops.EagerTensor):
+    if dtype is not None and value.dtype != dtype:
+      raise TypeError(f"Expected tensor {value} with dtype {dtype!r}, but got "
+                      f"dtype {value.dtype!r}.")
+    return value
+  if dtype is not None:
+    try:
+      dtype = dtype.as_datatype_enum
+    except AttributeError:
+      dtype = dtypes.as_dtype(dtype).as_datatype_enum
+  ctx.ensure_initialized()
+  return ops.EagerTensor(value, ctx.device_name, dtype)
+
+
+@tf_export(v1=["constant"])
+def constant_v1(
+    value, dtype=None, shape=None, name="Const", verify_shape=False
+) -> tensor_lib.Tensor:
+  """Creates a constant tensor.
+
+  The resulting tensor is populated with values of type `dtype`, as
+  specified by arguments `value` and (optionally) `shape` (see examples
+  below).
+
+  The argument `value` can be a constant value, or a list of values of type
+  `dtype`. If `value` is a list, then the length of the list must be less
+  than or equal to the number of elements implied by the `shape` argument (if
+  specified). In the case where the list length is less than the number of
+  elements specified by `shape`, the last element in the list will be used
+  to fill the remaining entries.
+
+  The argument `shape` is optional. If present, it specifies the dimensions of
+  the resulting tensor. If not present, the shape of `value` is used.
+
+  If the argument `dtype` is not specified, then the type is inferred from
+  the type of `value`.
+
+  For example:
+
+  ```python
+  # Constant 1-D Tensor populated with value list.
+  tensor = tf.constant([1, 2, 3, 4, 5, 6, 7]) => [1 2 3 4 5 6 7]
+
+  # Constant 2-D tensor populated with scalar value -1.
+  tensor = tf.constant(-1.0, shape=[2, 3]) => [[-1. -1. -1.]
+                                               [-1. -1. -1.]]
+  ```
+
+  `tf.constant` differs from `tf.fill` in a few ways:
+
+  *   `tf.constant` supports arbitrary constants, not just uniform scalar
+      Tensors like `tf.fill`.
+  *   `tf.constant` creates a `Const` node in the computation graph with the
+      exact value at graph construction time. On the other hand, `tf.fill`
+      creates an Op in the graph that is expanded at runtime.
+  *   Because `tf.constant` only embeds constant values in the graph, it does
+      not support dynamic shapes based on other runtime Tensors, whereas
+      `tf.fill` does.
+
+  Args:
+    value:          A constant value (or list) of output type `dtype`.
+
+    dtype:          The type of the elements of the resulting tensor.
+
+    shape:          Optional dimensions of resulting tensor.
+
+    name:           Optional name for the tensor.
+
+    verify_shape:   Boolean that enables verification of a shape of values.
+
+  Returns:
+    A Constant Tensor.
+
+  Raises:
+    TypeError: if shape is incorrectly specified or unsupported.
+  """
+  return _constant_impl(value, dtype, shape, name, verify_shape=verify_shape,
+                        allow_broadcast=False)
+
+
+@tf_export("constant", v1=[])
+def constant(value, dtype=None, shape=None, name="Const") -> tensor_lib.Tensor:
+  """Creates a constant tensor from a tensor-like object.
+
+  Note: All eager `tf.Tensor` values are immutable (in contrast to
+  `tf.Variable`). There is nothing especially _constant_ about the value
+  returned from `tf.constant`. This function is not fundamentally different from
+  `tf.convert_to_tensor`. The name `tf.constant` comes from the `value` being
+  embedded in a `Const` node in the `tf.Graph`. `tf.constant` is useful
+  for asserting that the value can be embedded that way.
+
+  If the argument `dtype` is not specified, then the type is inferred from
+  the type of `value`.
+
+  >>> # Constant 1-D Tensor from a python list.
+  >>> tf.constant([1, 2, 3, 4, 5, 6])
+  <tf.Tensor: shape=(6,), dtype=int32,
+      numpy=array([1, 2, 3, 4, 5, 6], dtype=int32)>
+  >>> # Or a numpy array
+  >>> a = np.array([[1, 2, 3], [4, 5, 6]])
+  >>> tf.constant(a)
+  <tf.Tensor: shape=(2, 3), dtype=int64, numpy=
+    array([[1, 2, 3],
+           [4, 5, 6]])>
+
+  If `dtype` is specified, the resulting tensor values are cast to the requested
+  `dtype`.
+
+  >>> tf.constant([1, 2, 3, 4, 5, 6], dtype=tf.float64)
+  <tf.Tensor: shape=(6,), dtype=float64,
+      numpy=array([1., 2., 3., 4., 5., 6.])>
+
+  If `shape` is set, the `value` is reshaped to match. Scalars are expanded to
+  fill the `shape`:
+
+  >>> tf.constant(0, shape=(2, 3))
+    <tf.Tensor: shape=(2, 3), dtype=int32, numpy=
+    array([[0, 0, 0],
+           [0, 0, 0]], dtype=int32)>
+  >>> tf.constant([1, 2, 3, 4, 5, 6], shape=[2, 3])
+  <tf.Tensor: shape=(2, 3), dtype=int32, numpy=
+    array([[1, 2, 3],
+           [4, 5, 6]], dtype=int32)>
+
+  `tf.constant` has no effect if an eager Tensor is passed as the `value`, it
+  even transmits gradients:
+
+  >>> v = tf.Variable([0.0])
+  >>> with tf.GradientTape() as g:
+  ...     loss = tf.constant(v + v)
+  >>> g.gradient(loss, v).numpy()
+  array([2.], dtype=float32)
+
+  But, since `tf.constant` embeds the value in the `tf.Graph` this fails for
+  symbolic tensors:
+
+  >>> with tf.compat.v1.Graph().as_default():
+  ...   i = tf.compat.v1.placeholder(shape=[None, None], dtype=tf.float32)
+  ...   t = tf.constant(i)
+  Traceback (most recent call last):
+  ...
+  TypeError: ...
+
+  `tf.constant` will create tensors on the current device. Inputs which are
+  already tensors maintain their placements unchanged.
+
+  Related Ops:
+
+  * `tf.convert_to_tensor` is similar but:
+    * It has no `shape` argument.
+    * Symbolic tensors are allowed to pass through.
+
+    >>> with tf.compat.v1.Graph().as_default():
+    ...   i = tf.compat.v1.placeholder(shape=[None, None], dtype=tf.float32)
+    ...   t = tf.convert_to_tensor(i)
+
+  * `tf.fill`: differs in a few ways:
+    *   `tf.constant` supports arbitrary constants, not just uniform scalar
+        Tensors like `tf.fill`.
+    *   `tf.fill` creates an Op in the graph that is expanded at runtime, so it
+        can efficiently represent large tensors.
+    *   Since `tf.fill` does not embed the value, it can produce dynamically
+        sized outputs.
+
+  Args:
+    value: A constant value (or list) of output type `dtype`.
+    dtype: Optional type of the elements of the resulting tensor.
+    shape: Optional dimensions of resulting tensor.
+    name: Optional name for the tensor.
+
+  Returns:
+    A Constant Tensor.
+
+  Raises:
+    TypeError: if shape is incorrectly specified or unsupported.
+    ValueError: if called on a symbolic tensor.
+  """
+  return _constant_impl(value, dtype, shape, name, verify_shape=False,
+                        allow_broadcast=True)
+
+
+def _constant_impl(
+    value, dtype, shape, name, verify_shape, allow_broadcast
+) -> tensor_lib.Tensor:
+  """Implementation of constant."""
+  ctx = context.context()
+  if ctx.executing_eagerly():
+    if trace.enabled:
+      with trace.Trace("tf.constant"):
+        return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
+    return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
+
+  const_tensor = ops._create_graph_constant(  # pylint: disable=protected-access
+      value, dtype, shape, name, verify_shape, allow_broadcast
+  )
+  return const_tensor
+
+
+def _constant_eager_impl(
+    ctx, value, dtype, shape, verify_shape
+) -> ops._EagerTensorBase:
+  """Creates a constant on the current device."""
+  t = convert_to_eager_tensor(value, ctx, dtype)
+  if shape is None:
+    return t
+  shape = tensor_shape.as_shape(shape)
+  if shape == t.shape:
+    return t
+  if verify_shape:
+    raise TypeError(f"Expected Tensor {t} (converted from {value}) with shape "
+                    f"{tuple(shape)}, but got shape {tuple(t.shape)}.")
+  num_t = t.shape.num_elements()
+  # TODO(josh11b): Implement shape -> eager tensor conversion.
+  if num_t == shape.num_elements():
+    return _eager_reshape(t, shape.as_list(), ctx)
+  if num_t == 1:
+    if t.dtype == dtypes.bool:
+      # We don't have a Fill kernel for bool dtype on GPU. So we first run
+      # Fill on CPU and then copy to GPU if needed.
+      with ops.device("/device:CPU:0"):
+        x = _eager_fill(shape.as_list(), _eager_identity(t, ctx), ctx)
+      return _eager_identity(x, ctx)
+    else:
+      return _eager_fill(shape.as_list(), t, ctx)
+  raise TypeError("Eager execution of tf.constant with unsupported shape. "
+                  f"Tensor {t} (converted from {value}) has {num_t:d} "
+                  f"elements, but got `shape` {shape} with "
+                  f"{shape.num_elements()} elements).")
+
+
+def is_constant(tensor_or_op):
+  if isinstance(tensor_or_op, tensor_lib.Tensor):
+    op = tensor_or_op.op
+  else:
+    op = tensor_or_op
+  return op.type == "Const"
+
+
+def _tensor_shape_tensor_conversion_function(s,
+                                             dtype=None,
+                                             name=None,
+                                             as_ref=False):
+  """Function to convert TensorShape to Tensor."""
+  _ = as_ref
+  if not s.is_fully_defined():
+    raise ValueError(
+        f"Cannot convert a partially known TensorShape {s} to a Tensor.")
+  s_list = s.as_list()
+  int64_value = 0
+  for dim in s_list:
+    if dim >= 2**31:
+      int64_value = dim
+      break
+
+  if dtype is not None:
+    if dtype not in (dtypes.int32, dtypes.int64):
+      raise TypeError(f"Cannot convert TensorShape {s} to dtype {dtype}. "
+                      "Allowed dtypes are tf.int32 and tf.int64.")
+    if dtype == dtypes.int32 and int64_value:
+      raise ValueError(f"Cannot convert TensorShape {s} to dtype int32; "
+                       f"a dimension is too large. Consider using tf.int64.")
+  else:
+    dtype = dtypes.int64 if int64_value else dtypes.int32
+  if name is None:
+    name = "shape_as_tensor"
+  return constant(s_list, dtype=dtype, name=name)
+
+
+tensor_conversion_registry.register_tensor_conversion_function(
+    tensor_shape.TensorShape, _tensor_shape_tensor_conversion_function, 100)
+
+
+def _dimension_tensor_conversion_function(d,
+                                          dtype=None,
+                                          name=None,
+                                          as_ref=False):
+  """Function to convert Dimension to Tensor."""
+  _ = as_ref
+  if d.value is None:
+    raise ValueError(f"Cannot convert unknown Dimension {d} to a Tensor.")
+  if dtype is not None:
+    if dtype not in (dtypes.int32, dtypes.int64):
+      raise TypeError(f"Cannot convert Dimension {d} to dtype {dtype}. "
+                      "Allowed dtypes are tf.int32 and tf.int64.")
+  else:
+    dtype = dtypes.int32
+  if name is None:
+    name = "shape_as_tensor"
+  return constant(d.value, dtype=dtype, name=name)
+
+
+tensor_conversion_registry.register_tensor_conversion_function(
+    tensor_shape.Dimension, _dimension_tensor_conversion_function, 100)
+
+
+class _ConstantTensorCodec:
+  """Codec for Tensor."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, tensor_lib.Tensor)
+
+  def do_encode(self, tensor_value, encode_fn):
+    """Returns an encoded `TensorProto` for the given `tf.Tensor`."""
+    del encode_fn
+    encoded_tensor = struct_pb2.StructuredValue()
+    if isinstance(tensor_value, ops.EagerTensor):
+      encoded_tensor.tensor_value.CopyFrom(
+          tensor_util.make_tensor_proto(tensor_value.numpy())
+      )
+    else:
+      if tensor_value.op.type == "Const":
+        encoded_tensor.tensor_value.CopyFrom(tensor_value.op.get_attr("value"))
+      else:
+        raise nested_structure_coder.NotEncodableError(
+            f"No encoder for object {str(tensor_value)} of type"
+            f" {type(tensor_value)}."
+        )
+    return encoded_tensor
+
+  def can_decode(self, value):
+    return value.HasField("tensor_value")
+
+  def do_decode(self, value, decode_fn):
+    """Returns the `tf.Tensor` encoded by the proto `value`."""
+    del decode_fn
+    tensor_proto = value.tensor_value
+    tensor = constant(tensor_util.MakeNdarray(tensor_proto))
+    return tensor
+
+
+nested_structure_coder.register_codec(_ConstantTensorCodec())
+
+
+class _NumpyCodec:
+  """Codec for Numpy."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, np.ndarray)
+
+  def do_encode(self, numpy_value, encode_fn):
+    """Returns an encoded `TensorProto` for `np.ndarray`."""
+    del encode_fn
+    encoded_numpy = struct_pb2.StructuredValue()
+    encoded_numpy.numpy_value.CopyFrom(
+        tensor_util.make_tensor_proto(numpy_value)
+    )
+    return encoded_numpy
+
+  def can_decode(self, value):
+    return value.HasField("numpy_value")
+
+  def do_decode(self, value, decode_fn):
+    """Returns the `np.ndarray` encoded by the proto `value`."""
+    del decode_fn
+    tensor_proto = value.numpy_value
+    numpy = tensor_util.MakeNdarray(tensor_proto)
+    return numpy
+
+
+nested_structure_coder.register_codec(_NumpyCodec())

tensorflow/python/framework/constant_op_test.py

@@ -0,0 +1,144 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.constant_op."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from google.protobuf import text_format
+
+from tensorflow.core.framework import graph_pb2
+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 importer
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops.parallel_for import control_flow_ops
+from tensorflow.python.platform import test
+
+
+class ConstantOpTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      dtypes.bfloat16,
+      dtypes.complex128,
+      dtypes.complex64,
+      dtypes.double,
+      dtypes.float16,
+      dtypes.float32,
+      dtypes.float64,
+      dtypes.half,
+      dtypes.int16,
+      dtypes.int32,
+      dtypes.int64,
+      dtypes.int8,
+      dtypes.qint16,
+      dtypes.qint32,
+      dtypes.qint8,
+      dtypes.quint16,
+      dtypes.quint8,
+      dtypes.uint16,
+      dtypes.uint32,
+      dtypes.uint64,
+      dtypes.uint8,
+  )
+  def test_convert_string_to_number(self, dtype):
+    with self.assertRaises(TypeError):
+      constant_op.constant("hello", dtype)
+
+  def _make_graph_def(self, text):
+    ret = graph_pb2.GraphDef()
+    text_format.Parse(text, ret)
+    return ret
+
+  def test_eager_const_xla(self):
+
+    @def_function.function(jit_compile=True)
+    def f_using_eagerconst(x):
+      graph_def = self._make_graph_def("""
+         node { name: 'x' op: 'Const'
+           attr { key: 'dtype' value { type: DT_FLOAT } }
+           attr { key: 'value' value { tensor {
+             dtype: DT_FLOAT tensor_shape {} float_val: NaN } } } }
+         node { name: 'const' op: '_EagerConst' input: 'x:0'
+                attr { key: 'T' value { type: DT_FLOAT } }}""")
+      x_id = importer.import_graph_def(
+          graph_def,
+          input_map={"x:0": x},
+          return_elements=["const"],
+          name="import")[0].outputs[0]
+      return x_id
+
+    self.assertAllClose(3.14, f_using_eagerconst(constant_op.constant(3.14)))
+
+  def test_np_array_memory_not_shared(self):
+    # An arbitrarily large loop number to test memory sharing
+    for _ in range(10000):
+      x = np.arange(10)
+      xt = constant_op.constant(x)
+      x[3] = 42
+      # Changing the input array after `xt` is created should not affect `xt`
+      self.assertEqual(xt.numpy()[3], 3)
+
+  def test_eager_const_grad_error(self):
+
+    @def_function.function
+    def f_using_eagerconst():
+      x = constant_op.constant(1.)
+      graph_def = self._make_graph_def("""
+         node { name: 'x' op: 'Placeholder'
+                attr { key: 'dtype' value { type: DT_FLOAT } }}
+         node { name: 'const' op: '_EagerConst' input: 'x:0'
+                attr { key: 'T' value { type: DT_FLOAT } }}""")
+      x_id = importer.import_graph_def(
+          graph_def,
+          input_map={"x:0": x},
+          return_elements=["const"],
+          name="import")[0].outputs[0]
+      gradients_impl.gradients(x_id, x)
+      return x_id
+
+    with self.assertRaisesRegex(AssertionError, "Please file a bug"):
+      f_using_eagerconst()
+
+  def test_eager_const_pfor(self):
+
+    @def_function.function
+    def f_using_eagerconst():
+
+      def vec_fn(x):
+        graph_def = self._make_graph_def("""
+           node { name: 'x' op: 'Const'
+             attr { key: 'dtype' value { type: DT_FLOAT } }
+             attr { key: 'value' value { tensor {
+               dtype: DT_FLOAT tensor_shape {} float_val: 3.14 } } } }
+           node { name: 'const' op: '_EagerConst' input: 'x:0'
+                  attr { key: 'T' value { type: DT_FLOAT } }}""")
+        return importer.import_graph_def(
+            graph_def,
+            input_map={"x:0": x},
+            return_elements=["const"],
+            name="import")[0].outputs[0]
+
+      return control_flow_ops.vectorized_map(
+          vec_fn, constant_op.constant([1., 2.]), fallback_to_while_loop=False)
+
+    self.assertAllClose([1., 2.], f_using_eagerconst())
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/framework/constant_tensor_conversion.py

@@ -0,0 +1,45 @@
+# 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.
+# ==============================================================================
+"""Tensor conversion factory functions for builtins to constant Tensors."""
+
+from tensorflow.python.framework import tensor_conversion_registry
+
+
+# Factory function for tensor conversion for builtins. Import constant_op.py
+# in-line so that it is only imported when it is needed. This file is imported
+# at TF import time, thus that helps reduce import slowness.
+def _constant_tensor_conversion_function(
+    v, dtype=None, name=None, as_ref=False
+):
+  from tensorflow.python.framework import constant_op  # pylint: disable=g-import-not-at-top
+
+  _ = as_ref
+  return constant_op.constant(v, dtype=dtype, name=name)
+
+
+# Register the conversion function for the "unconvertible" types
+# as a conversion to a constant.
+tensor_conversion_registry.register_tensor_conversion_function_internal(
+    tensor_conversion_registry._CONSTANT_OP_CONVERTIBLES,  # pylint: disable=protected-access
+    _constant_tensor_conversion_function,
+    0,
+)
+
+tensor_conversion_registry.register_tensor_conversion_function(
+    (list, tuple), _constant_tensor_conversion_function, 100
+)
+tensor_conversion_registry.register_tensor_conversion_function(
+    object, _constant_tensor_conversion_function, 200
+)

tensorflow/python/framework/cpp_shape_inference.proto

@@ -0,0 +1,53 @@
+// 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 = "proto3";
+
+package tensorflow;
+
+import "tensorflow/core/framework/full_type.proto";
+import "tensorflow/core/framework/tensor_shape.proto";
+import "tensorflow/core/framework/types.proto";
+
+option cc_enable_arenas = true;
+option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/python/framework/cpp_shape_inference_go_proto";
+
+// LINT.IfChange
+message CppShapeInferenceResult {
+  message HandleShapeAndType {
+    reserved 3;
+
+    TensorShapeProto shape = 1;
+    DataType dtype = 2;
+    FullTypeDef type = 4;
+  }
+  message HandleData {
+    bool is_set = 1;
+
+    // Only valid if <is_set>.
+    repeated HandleShapeAndType shape_and_type = 2;
+  }
+  TensorShapeProto shape = 1;
+
+  reserved 2;  // was handle_shape
+  reserved 3;  // was handle_dtype
+  HandleData handle_data = 4;
+}
+
+message CppShapeInferenceInputsNeeded {
+  repeated int32 input_tensors_needed = 1;
+  repeated int32 input_tensors_as_shapes_needed = 2;
+}
+// LINT.ThenChange(//tensorflow/core/framework/cpp_shape_inference.proto)

tensorflow/python/framework/device.py

@@ -0,0 +1,178 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Class to represent a device."""
+
+from tensorflow.python import tf2
+from tensorflow.python.framework import device_spec
+
+if tf2.enabled():
+  DeviceSpec = device_spec.DeviceSpecV2
+else:
+  DeviceSpec = device_spec.DeviceSpecV1
+
+
+def check_valid(spec):
+  """Check that a device spec is valid.
+
+  Args:
+    spec: a string.
+
+  Raises:
+    An exception if the spec is invalid.
+  """
+  # Construct a DeviceSpec.  It will assert a failure if spec is invalid.
+  DeviceSpec.from_string(spec)
+
+
+def is_device_spec(obj):
+  """Abstract away the fact that DeviceSpecV2 is the base class."""
+  return isinstance(obj, device_spec.DeviceSpecV2)
+
+
+def canonical_name(device):
+  """Returns a canonical name for the given `DeviceSpec` or device name."""
+  if device is None:
+    return ""
+  if is_device_spec(device):
+    return device.to_string()
+  else:
+    device = DeviceSpec.from_string(device)
+    return device.to_string()
+
+
+# Performance caches
+_cached_mergers = {}
+_string_merge_cache = {}
+
+
+def merge_device(spec):
+  """Returns a device function that merges devices specifications.
+
+  This can be used to merge partial specifications of devices. The
+  innermost setting for a device field takes precedence. For example:
+
+    with tf.device(merge_device("/device:GPU:0"))
+      # Nodes created here have device "/device:GPU:0"
+      with tf.device(merge_device("/job:worker")):
+        # Nodes created here have device "/job:worker/device:GPU:0"
+        with tf.device(merge_device("/device:CPU:0")):
+          # Nodes created here have device "/job:worker/device:CPU:0"
+          with tf.device(merge_device("/job:ps")):
+            # Nodes created here have device "/job:ps/device:CPU:0"
+
+  Args:
+    spec: A `DeviceSpec` or a device spec string (partially) describing the
+      device that should be used for all nodes created in the scope of
+      the returned device function's with block.
+
+  Returns:
+    A MergeDevice object with the above-described behavior.
+
+  Raises:
+    ValueError: if the spec was not valid.
+  """
+
+  if isinstance(spec, MergeDevice):
+    return spec
+
+  merger = _cached_mergers.get(spec)
+  if merger:
+    return merger
+  merger = MergeDevice(spec)
+  # No locking needed, since updates are stateless.
+  _cached_mergers[spec] = merger
+  return merger
+
+
+class MergeDevice(object):
+  """Wraps a device specification (DeviceSpec or str) with merge functionality.
+
+  When called, this class will merge a node_def with its own spec. It also
+  exposes a `shortcut_string_merge` method which can significantly improve
+  performance of device placement.
+  """
+
+  __slots__ = ["_spec"]
+
+  def __init__(self, spec):
+    if isinstance(spec, device_spec.DeviceSpecV2):
+      self._spec = spec
+    elif isinstance(spec, device_spec.DeviceSpecV1):
+      # Capture a snapshot of spec.
+      self._spec = spec.__class__.from_string(spec.to_string())
+    else:
+      self._spec = DeviceSpec.from_string(spec)
+
+  def __call__(self, node_def):
+    # In general a user may create a device function which takes into account
+    # arbitrary properties of an op. (For instance dynamically placing ops based
+    # on type.) So even though the standard DeviceSpec route only uses the
+    # device attribute, we take an entire node_def to maintain a consistent
+    # signature with general device functions.
+    current_device = DeviceSpec.from_string(node_def.device or "")
+    return self._spec.make_merged_spec(current_device)
+
+  def shortcut_string_merge(self, node_def):
+    """Merge a node def without materializing a full DeviceSpec object.
+
+    Often a device merge is invoked in order to generate a string which can be
+    passed into the c api. In such a case, we can cache the
+      node_def.device  ->  merge_result_string
+
+    map, and in most cases avoid:
+      - Materializing a copy of self._spec (In the case of DeviceSpecV1)
+      - Materializing a DeviceSpec for node_def.device
+      - A DeviceSpec.merge_from invocation
+
+    In practice the cache hit rate for this function is very high, because the
+    number of invocations when iterating through the device stack is much
+    larger than the number of devices.
+
+    Args:
+      node_def: An Operation (or Operation-like) to merge device constraints
+        with self._spec
+
+    Returns:
+      A string containing the merged device specification.
+    """
+    device = node_def.device or ""
+
+    merge_key = (self._spec, device)
+    result = _string_merge_cache.get(merge_key)
+    if result is None:
+      # This update is not atomic, however because the merge is stateless
+      # we don't need to lock when updating the cache.
+      result = self.__call__(node_def).to_string()
+      _string_merge_cache[merge_key] = result
+
+    return result
+
+  def __repr__(self):
+    return "{} (spec: {})".format(
+        super(MergeDevice, self).__repr__(), self._spec.to_string())
+
+  @property
+  def is_null_merge(self):
+    """Indicate whether the wrapped spec is empty.
+
+    In the degenerate case where self._spec is an empty specification, a caller
+    may wish to skip a merge step entirely. (However this class does not have
+    enough information to make that determination.)
+
+    Returns:
+      A boolean indicating whether a device merge will be trivial.
+    """
+    return not bool(self._spec.to_string())

tensorflow/python/framework/device_spec.py

@@ -0,0 +1,485 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Class to represent a device."""
+
+from tensorflow.python.util.tf_export import tf_export
+from tensorflow.python import pywrap_tfe
+
+# EPU represents for TPU embedding for now. Subject to change in future.
+_VALID_DEVICE_TYPES = frozenset({"CPU", "GPU", "TPU", "CUSTOM", "EPU"})
+
+# ==============================================================================
+# == Global Implementation Details =============================================
+# ==============================================================================
+_STRING_TO_COMPONENTS_CACHE = {}
+_COMPONENTS_TO_STRING_CACHE = {}
+
+
+def _as_str_or_none(inp):
+  return None if inp is None else str(inp)
+
+
+def _as_int_or_none(inp):
+  return None if inp is None else int(inp)
+
+
+def _as_device_str_or_none(device_type):
+  # For backwards compatibility only, we support lowercase variants of
+  # cpu and gpu but turn them into uppercase here.
+  if device_type in ("cpu", "gpu"):
+    return device_type.upper()
+  return _as_str_or_none(device_type)
+
+
+@tf_export("DeviceSpec", v1=[])
+class DeviceSpecV2(object):
+  """Represents a (possibly partial) specification for a TensorFlow device.
+
+  `DeviceSpec`s are used throughout TensorFlow to describe where state is stored
+  and computations occur. Using `DeviceSpec` allows you to parse device spec
+  strings to verify their validity, merge them or compose them programmatically.
+
+  Example:
+
+  ```python
+  # Place the operations on device "GPU:0" in the "ps" job.
+  device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
+  with tf.device(device_spec.to_string()):
+    # Both my_var and squared_var will be placed on /job:ps/device:GPU:0.
+    my_var = tf.Variable(..., name="my_variable")
+    squared_var = tf.square(my_var)
+  ```
+
+  With eager execution disabled (by default in TensorFlow 1.x and by calling
+  disable_eager_execution() in TensorFlow 2.x), the following syntax
+  can be used:
+
+  ```python
+  tf.compat.v1.disable_eager_execution()
+
+  # Same as previous
+  device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
+  # No need of .to_string() method.
+  with tf.device(device_spec):
+    my_var = tf.Variable(..., name="my_variable")
+    squared_var = tf.square(my_var)
+  ```
+
+  If a `DeviceSpec` is partially specified, it will be merged with other
+  `DeviceSpec`s according to the scope in which it is defined. `DeviceSpec`
+  components defined in inner scopes take precedence over those defined in
+  outer scopes.
+
+  ```python
+  gpu0_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)
+  with tf.device(DeviceSpec(job="train").to_string()):
+    with tf.device(gpu0_spec.to_string()):
+      # Nodes created here will be assigned to /job:ps/device:GPU:0.
+    with tf.device(DeviceSpec(device_type="GPU", device_index=1).to_string()):
+      # Nodes created here will be assigned to /job:train/device:GPU:1.
+  ```
+
+  A `DeviceSpec` consists of 5 components -- each of
+  which is optionally specified:
+
+  * Job: The job name.
+  * Replica: The replica index.
+  * Task: The task index.
+  * Device type: The device type string (e.g. "CPU" or "GPU").
+  * Device index: The device index.
+  """
+
+  __slots__ = ("_job", "_replica", "_task", "_device_type", "_device_index",
+               "_as_string", "_hash")
+
+  def __init__(self,
+               job=None,
+               replica=None,
+               task=None,
+               device_type=None,
+               device_index=None):
+    """Create a new `DeviceSpec` object.
+
+    Args:
+      job: string.  Optional job name.
+      replica: int.  Optional replica index.
+      task: int.  Optional task index.
+      device_type: Optional device type string (e.g. "CPU" or "GPU")
+      device_index: int.  Optional device index.  If left unspecified, device
+        represents 'any' device_index.
+    """
+    self._job = _as_str_or_none(job)
+    self._replica = _as_int_or_none(replica)
+    self._task = _as_int_or_none(task)
+    self._device_type = _as_device_str_or_none(device_type)
+    self._device_index = _as_int_or_none(device_index)
+    self._as_string = self._components_to_string(
+        job=self._job,
+        replica=self._replica,
+        task=self._task,
+        device_type=self._device_type,
+        device_index=self._device_index)
+    self._hash = hash(self.to_string())
+
+  def to_string(self):
+    """Return a string representation of this `DeviceSpec`.
+
+    Returns:
+      a string of the form
+      /job:<name>/replica:<id>/task:<id>/device:<device_type>:<id>.
+    """
+    return self._as_string
+
+  @classmethod
+  def from_string(cls, spec):
+    """Construct a `DeviceSpec` from a string.
+
+    Args:
+      spec: a string of the form
+       /job:<name>/replica:<id>/task:<id>/device:CPU:<id> or
+       /job:<name>/replica:<id>/task:<id>/device:GPU:<id> as cpu and gpu are
+         mutually exclusive. All entries are optional.
+
+    Returns:
+      A DeviceSpec.
+    """
+    return cls(*cls._string_to_components(spec))
+
+  def parse_from_string(self, spec):
+    """Parse a `DeviceSpec` name into its components.
+
+    **2.x behavior change**:
+
+    In TensorFlow 1.x, this function mutates its own state and returns itself.
+    In 2.x, DeviceSpecs are immutable, and this function will return a
+      DeviceSpec which contains the spec.
+
+    * Recommended:
+
+      ```
+      # my_spec and my_updated_spec are unrelated.
+      my_spec = tf.DeviceSpec.from_string("/CPU:0")
+      my_updated_spec = tf.DeviceSpec.from_string("/GPU:0")
+      with tf.device(my_updated_spec):
+        ...
+      ```
+
+    * Will work in 1.x and 2.x (though deprecated in 2.x):
+
+      ```
+      my_spec = tf.DeviceSpec.from_string("/CPU:0")
+      my_updated_spec = my_spec.parse_from_string("/GPU:0")
+      with tf.device(my_updated_spec):
+        ...
+      ```
+
+    * Will NOT work in 2.x:
+
+      ```
+      my_spec = tf.DeviceSpec.from_string("/CPU:0")
+      my_spec.parse_from_string("/GPU:0")  # <== Will not update my_spec
+      with tf.device(my_spec):
+        ...
+      ```
+
+    In general, `DeviceSpec.from_string` should completely replace
+    `DeviceSpec.parse_from_string`, and `DeviceSpec.replace` should
+    completely replace setting attributes directly.
+
+    Args:
+      spec: an optional string of the form
+       /job:<name>/replica:<id>/task:<id>/device:CPU:<id> or
+       /job:<name>/replica:<id>/task:<id>/device:GPU:<id> as cpu and gpu are
+         mutually exclusive. All entries are optional.
+
+    Returns:
+      The `DeviceSpec`.
+
+    Raises:
+      ValueError: if the spec was not valid.
+    """
+    return self.from_string(spec)
+
+  def make_merged_spec(self, dev):
+    """Returns a new DeviceSpec which incorporates `dev`.
+
+    When combining specs, `dev` will take precedence over the current spec.
+    So for instance:
+    ```
+    first_spec = tf.DeviceSpec(job=0, device_type="CPU")
+    second_spec = tf.DeviceSpec(device_type="GPU")
+    combined_spec = first_spec.make_merged_spec(second_spec)
+    ```
+
+    is equivalent to:
+    ```
+    combined_spec = tf.DeviceSpec(job=0, device_type="GPU")
+    ```
+
+    Args:
+      dev: a `DeviceSpec`
+
+    Returns:
+      A new `DeviceSpec` which combines `self` and `dev`
+    """
+    return self.__class__(*self._get_combined_properties(dev))
+
+  def replace(self, **kwargs):
+    """Convenience method for making a new DeviceSpec by overriding fields.
+
+    For instance:
+    ```
+    my_spec = DeviceSpec=(job="my_job", device="CPU")
+    my_updated_spec = my_spec.replace(device="GPU")
+    my_other_spec = my_spec.replace(device=None)
+    ```
+
+    Args:
+      **kwargs: This method takes the same args as the DeviceSpec constructor
+
+    Returns:
+      A DeviceSpec with the fields specified in kwargs overridden.
+    """
+    init_kwargs = dict(
+        job=self.job,
+        replica=self.replica,
+        task=self.task,
+        device_type=self.device_type,
+        device_index=self.device_index)
+
+    # Explicitly provided kwargs take precedence.
+    init_kwargs.update(kwargs)
+    return self.__class__(**init_kwargs)
+
+  @property
+  def job(self):
+    return self._job
+
+  @property
+  def replica(self):
+    return self._replica
+
+  @property
+  def task(self):
+    return self._task
+
+  @property
+  def device_type(self):
+    return self._device_type
+
+  @property
+  def device_index(self):
+    return self._device_index
+
+  def _get_combined_properties(self, dev):
+    """Combine the current DeviceSpec with another DeviceSpec.
+
+    The combination of DeviceSpecs is will give priority to dev.
+
+    Args:
+      dev: a `DeviceSpec`
+
+    Returns:
+      A tuple of (job, replica, task, device_type, device_index) which
+      represents the combination of self and dev.
+    """
+    return (
+        dev.job if dev.job is not None else self.job,
+        dev.replica if dev.replica is not None else self.replica,
+        dev.task if dev.task is not None else self.task,
+        dev.device_type if dev.device_type is not None else self.device_type,
+        dev.device_index if dev.device_index is not None else self.device_index,
+    )
+
+  @staticmethod
+  def _get_valid_device_types():
+    valid_device_types = set({})
+    physical_devices = pywrap_tfe.TF_ListPluggablePhysicalDevices()
+    for device in physical_devices:
+      valid_device_types.add(device.decode().split(":")[1])
+    valid_device_types = valid_device_types | _VALID_DEVICE_TYPES
+    return valid_device_types
+
+  @staticmethod
+  def _string_to_components(spec=None):
+    """Stateless portion of device spec string parsing.
+
+    Args:
+      spec: An optional string specifying a device specification.
+
+    Returns:
+      The parsed components of `spec`. Note that the result of this function
+      must go through attribute setters of DeviceSpec, and should therefore NOT
+      be used directly.
+    """
+    cached_result = _STRING_TO_COMPONENTS_CACHE.get(spec)
+    if cached_result is not None:
+      return cached_result
+
+    raw_spec = spec  # keep a copy of the original to update the cache
+    job, replica, task, device_type, device_index = None, None, None, None, None
+
+    spec = spec or ""
+    splits = [x.split(":") for x in spec.split("/")]
+    valid_device_types = DeviceSpecV2._get_valid_device_types()
+    for y in splits:
+      ly = len(y)
+      if y:
+        # NOTE(taylorrobie): these will go through setters later.
+        if ly == 2 and y[0] == "job":
+          job = y[1]
+        elif ly == 2 and y[0] == "replica":
+          replica = y[1]
+        elif ly == 2 and y[0] == "task":
+          task = y[1]
+        elif ((ly == 1 or ly == 2) and (y[0].upper() in valid_device_types)):
+          if device_type is not None:
+            raise ValueError(f"Multiple device types are not allowed "
+                             f"while parsing the device spec: {spec}.")
+          device_type = y[0].upper()
+          if ly == 2 and y[1] != "*":
+            device_index = int(y[1])
+        elif ly == 3 and y[0] == "device":
+          if device_type is not None:
+            raise ValueError(f"Multiple device types are not allowed "
+                             f"while parsing the device spec: {spec}.")
+          device_type = y[1]
+          if y[2] != "*":
+            device_index = int(y[2])
+        elif ly and y[0] != "":  # pylint: disable=g-explicit-bool-comparison
+          raise ValueError(f"Unknown attribute '{y[0]}' is encountered "
+                           f"while parsing the device spec: '{spec}'.")
+
+    output = (job, replica, task, device_type, device_index)
+    _STRING_TO_COMPONENTS_CACHE[raw_spec] = output
+    return output
+
+  @staticmethod
+  def _components_to_string(job, replica, task, device_type, device_index):
+    """Stateless portion of `to_string` (separated to allow caching)."""
+    key = (job, replica, task, device_type, device_index)
+    cached_result = _COMPONENTS_TO_STRING_CACHE.get(key)
+    if cached_result is not None:
+      return cached_result
+
+    output = []
+    if job is not None:
+      output.append("/job:" + job)
+    if replica is not None:
+      output.append("/replica:" + str(replica))
+    if task is not None:
+      output.append("/task:" + str(task))
+    if device_type is not None:
+      device_index_string = "*"
+      if device_index is not None:
+        # Unlike the others, device_index is stored as an int.
+        device_index_string = str(device_index)
+      output.append("/device:%s:%s" % (device_type, device_index_string))
+
+    output = "".join(output)
+    _COMPONENTS_TO_STRING_CACHE[key] = output
+    return output
+
+  def __eq__(self, other):
+    """Checks if the `other` DeviceSpec is same as the current instance, eg have
+
+       same value for all the internal fields.
+
+    Args:
+      other: Another DeviceSpec
+
+    Returns:
+      Return `True` if `other` is also a DeviceSpec instance and has same value
+      as the current instance.
+      Return `False` otherwise.
+    """
+    return (isinstance(other, self.__class__) and
+            self.to_string() == other.to_string())
+
+  def __hash__(self):
+    return self._hash
+
+  def __repr__(self):
+    return (
+        f"<DeviceSpec(job={self.job}, replica={self.replica}, task={self.task}, "
+        f"device_type={self.device_type}, device_index={self.device_index})>")
+
+
+@tf_export(v1=["DeviceSpec"])  # pylint: disable=missing-docstring
+class DeviceSpecV1(DeviceSpecV2):
+  __doc__ = DeviceSpecV2.__doc__
+  __slots__ = DeviceSpecV2.__slots__
+
+  @DeviceSpecV2.job.setter
+  def job(self, job):
+    self._job = _as_str_or_none(job)
+    self._as_string, self._hash = None, None
+
+  @DeviceSpecV2.replica.setter
+  def replica(self, replica):
+    self._replica = _as_int_or_none(replica)
+    self._as_string, self._hash = None, None
+
+  @DeviceSpecV2.task.setter
+  def task(self, task):
+    self._task = _as_int_or_none(task)
+    self._as_string, self._hash = None, None
+
+  @DeviceSpecV2.device_type.setter
+  def device_type(self, device_type):
+    self._device_type = _as_device_str_or_none(device_type)
+    self._as_string, self._hash = None, None
+
+  @DeviceSpecV2.device_index.setter
+  def device_index(self, device_index):
+    self._device_index = _as_int_or_none(device_index)
+    self._as_string, self._hash = None, None
+
+  def __hash__(self):
+    if self._hash is None:
+      self._hash = hash(self.to_string())
+    return self._hash
+
+  def to_string(self):
+    if self._as_string is None:
+      self._as_string = self._components_to_string(
+          job=self.job,
+          replica=self.replica,
+          task=self.task,
+          device_type=self.device_type,
+          device_index=self.device_index)
+    return self._as_string
+
+  def parse_from_string(self, spec):
+    (self.job, self.replica, self.task, self.device_type,
+     self.device_index) = self._string_to_components(spec)
+
+    return self
+
+  def merge_from(self, dev):
+    """Merge the properties of "dev" into this `DeviceSpec`.
+
+    Note: Will be removed in TensorFlow 2.x since DeviceSpecs will become
+          immutable.
+
+    Args:
+      dev: a `DeviceSpec`.
+    """
+    (self.job, self.replica, self.task, self.device_type,
+     self.device_index) = self._get_combined_properties(dev)
+
+  # Use parent class docstrings for public methods.
+  to_string.__doc__ = DeviceSpecV2.to_string.__doc__
+  parse_from_string.__doc__ = DeviceSpecV2.parse_from_string.__doc__

tensorflow/python/framework/device_spec_test.py

@@ -0,0 +1,227 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.device_spec."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import device_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+
+TEST_V1_AND_V2 = (("v1", device_spec.DeviceSpecV1),
+                  ("v2", device_spec.DeviceSpecV2))
+
+
+class DeviceSpecTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def test_empty(self, device_spec_type):
+    d = device_spec_type()
+    self.assertEqual("", d.to_string())
+    d.parse_from_string("")
+    self.assertEqual("", d.to_string())
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def test_constructor(self, device_spec_type):
+    d = device_spec_type(job="j", replica=0, task=1,
+                         device_type="CPU", device_index=2)
+    self.assertEqual("j", d.job)
+    self.assertEqual(0, d.replica)
+    self.assertEqual(1, d.task)
+    self.assertEqual("CPU", d.device_type)
+    self.assertEqual(2, d.device_index)
+    self.assertEqual("/job:j/replica:0/task:1/device:CPU:2", d.to_string())
+
+    d = device_spec_type(device_type="GPU", device_index=0)
+    self.assertEqual("/device:GPU:0", d.to_string())
+
+  def testto_string_legacy(self):
+    """DeviceSpecV1 allows direct mutation."""
+    d = device_spec.DeviceSpecV1()
+    d.job = "foo"
+    self.assertEqual("/job:foo", d.to_string())
+    d.task = 3
+    self.assertEqual("/job:foo/task:3", d.to_string())
+    d.device_type = "CPU"
+    d.device_index = 0
+    self.assertEqual("/job:foo/task:3/device:CPU:0", d.to_string())
+    d.task = None
+    d.replica = 12
+    self.assertEqual("/job:foo/replica:12/device:CPU:0", d.to_string())
+    d.device_type = "GPU"
+    d.device_index = 2
+    self.assertEqual("/job:foo/replica:12/device:GPU:2", d.to_string())
+    d.device_type = "CPU"
+    d.device_index = 1
+    self.assertEqual("/job:foo/replica:12/device:CPU:1", d.to_string())
+    d.device_type = None
+    d.device_index = None
+    self.assertEqual("/job:foo/replica:12", d.to_string())
+
+    # Test wildcard
+    d = device_spec.DeviceSpecV1(job="foo", replica=12, task=3,
+                                 device_type="GPU")
+    self.assertEqual("/job:foo/replica:12/task:3/device:GPU:*", d.to_string())
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def test_replace(self, device_spec_type):
+    d = device_spec_type()
+    d = d.replace(job="foo")
+    self.assertEqual("/job:foo", d.to_string())
+
+    d = d.replace(task=3)
+    self.assertEqual("/job:foo/task:3", d.to_string())
+
+    d = d.replace(device_type="CPU", device_index=0)
+    self.assertEqual("/job:foo/task:3/device:CPU:0", d.to_string())
+
+    d = d.replace(task=None, replica=12)
+    self.assertEqual("/job:foo/replica:12/device:CPU:0", d.to_string())
+
+    d = d.replace(device_type="GPU", device_index=2)
+    self.assertEqual("/job:foo/replica:12/device:GPU:2", d.to_string())
+
+    d = d.replace(device_type="CPU", device_index=1)
+    self.assertEqual("/job:foo/replica:12/device:CPU:1", d.to_string())
+
+    d = d.replace(device_type=None, device_index=None)
+    self.assertEqual("/job:foo/replica:12", d.to_string())
+
+    # Test wildcard
+    d = device_spec.DeviceSpecV1(job="foo", replica=12, task=3,
+                                 device_type="GPU")
+    self.assertEqual("/job:foo/replica:12/task:3/device:GPU:*", d.to_string())
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def testto_string(self, device_spec_type):
+    d = device_spec_type(job="foo")
+    self.assertEqual("/job:foo", d.to_string())
+
+    d = device_spec_type(job="foo", task=3)
+    self.assertEqual("/job:foo/task:3", d.to_string())
+
+    d = device_spec_type(job="foo", task=3, device_type="cpu", device_index=0)
+    self.assertEqual("/job:foo/task:3/device:CPU:0", d.to_string())
+
+    d = device_spec_type(job="foo", replica=12, device_type="cpu",
+                         device_index=0)
+    self.assertEqual("/job:foo/replica:12/device:CPU:0", d.to_string())
+
+    d = device_spec_type(job="foo", replica=12, device_type="gpu",
+                         device_index=2)
+    self.assertEqual("/job:foo/replica:12/device:GPU:2", d.to_string())
+
+    d = device_spec_type(job="foo", replica=12)
+    self.assertEqual("/job:foo/replica:12", d.to_string())
+
+    # Test wildcard
+    d = device_spec_type(job="foo", replica=12, task=3, device_type="GPU")
+    self.assertEqual("/job:foo/replica:12/task:3/device:GPU:*", d.to_string())
+
+  def test_parse_legacy(self):
+    d = device_spec.DeviceSpecV1()
+    d.parse_from_string("/job:foo/replica:0")
+    self.assertEqual("/job:foo/replica:0", d.to_string())
+
+    d.parse_from_string("/replica:1/task:0/cpu:0")
+    self.assertEqual("/replica:1/task:0/device:CPU:0", d.to_string())
+
+    d.parse_from_string("/replica:1/task:0/device:CPU:0")
+    self.assertEqual("/replica:1/task:0/device:CPU:0", d.to_string())
+
+    d.parse_from_string("/job:muu/device:GPU:2")
+    self.assertEqual("/job:muu/device:GPU:2", d.to_string())
+
+    with self.assertRaisesRegex(ValueError,
+                                "Multiple device types are not allowed"):
+      d.parse_from_string("/job:muu/device:GPU:2/cpu:0")
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def test_to_from_string(self, device_spec_type):
+    d = device_spec_type.from_string("/job:foo/replica:0")
+    self.assertEqual("/job:foo/replica:0", d.to_string())
+    self.assertEqual(0, d.replica)
+
+    d = device_spec_type.from_string("/replica:1/task:0/cpu:0")
+    self.assertEqual("/replica:1/task:0/device:CPU:0", d.to_string())
+    self.assertAllEqual([1, 0, "CPU", 0],
+                        [d.replica, d.task, d.device_type, d.device_index])
+
+    d = device_spec_type.from_string("/replica:1/task:0/device:CPU:0")
+    self.assertEqual("/replica:1/task:0/device:CPU:0", d.to_string())
+    self.assertAllEqual([1, 0, "CPU", 0],
+                        [d.replica, d.task, d.device_type, d.device_index])
+
+    d = device_spec_type.from_string("/job:muu/device:GPU:2")
+    self.assertEqual("/job:muu/device:GPU:2", d.to_string())
+    self.assertAllEqual(["muu", "GPU", 2],
+                        [d.job, d.device_type, d.device_index])
+
+    with self.assertRaisesRegex(ValueError,
+                                "Multiple device types are not allowed"):
+      d.parse_from_string("/job:muu/device:GPU:2/cpu:0")
+
+  def test_merge_legacy(self):
+    d = device_spec.DeviceSpecV1.from_string("/job:foo/replica:0")
+    self.assertEqual("/job:foo/replica:0", d.to_string())
+
+    d.merge_from(device_spec.DeviceSpecV1.from_string("/task:1/device:GPU:2"))
+    self.assertEqual("/job:foo/replica:0/task:1/device:GPU:2", d.to_string())
+
+    d = device_spec.DeviceSpecV1()
+    d.merge_from(device_spec.DeviceSpecV1.from_string("/task:1/cpu:0"))
+    self.assertEqual("/task:1/device:CPU:0", d.to_string())
+
+    d.merge_from(device_spec.DeviceSpecV1.from_string("/job:boo/device:GPU:0"))
+    self.assertEqual("/job:boo/task:1/device:GPU:0", d.to_string())
+
+    d.merge_from(device_spec.DeviceSpecV1.from_string("/job:muu/cpu:2"))
+    self.assertEqual("/job:muu/task:1/device:CPU:2", d.to_string())
+    d.merge_from(device_spec.DeviceSpecV1.from_string(
+        "/job:muu/device:MyFunnyDevice:2"))
+    self.assertEqual("/job:muu/task:1/device:MyFunnyDevice:2", d.to_string())
+
+  def test_merge_removed(self):
+    with self.assertRaises(AttributeError):
+      d = device_spec.DeviceSpecV2()
+      d.merge_from(device_spec.DeviceSpecV2.from_string("/task:1/cpu:0"))
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def test_combine(self, device_spec_type):
+    d = device_spec_type.from_string("/job:foo/replica:0")
+    self.assertEqual("/job:foo/replica:0", d.to_string())
+
+    d = d.make_merged_spec(
+        device_spec_type.from_string("/task:1/device:GPU:2"))
+    self.assertEqual("/job:foo/replica:0/task:1/device:GPU:2", d.to_string())
+
+    d = device_spec_type()
+    d = d.make_merged_spec(device_spec_type.from_string("/task:1/cpu:0"))
+    self.assertEqual("/task:1/device:CPU:0", d.to_string())
+
+    d = d.make_merged_spec(
+        device_spec_type.from_string("/job:boo/device:GPU:0"))
+    self.assertEqual("/job:boo/task:1/device:GPU:0", d.to_string())
+
+    d = d.make_merged_spec(device_spec_type.from_string("/job:muu/cpu:2"))
+    self.assertEqual("/job:muu/task:1/device:CPU:2", d.to_string())
+    d = d.make_merged_spec(device_spec_type.from_string(
+        "/job:muu/device:MyFunnyDevice:2"))
+    self.assertEqual("/job:muu/task:1/device:MyFunnyDevice:2", d.to_string())
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/framework/device_test.py

@@ -0,0 +1,91 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.device."""
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import device
+from tensorflow.python.framework import device_spec
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import googletest
+
+TEST_V1_AND_V2 = (("v1", device_spec.DeviceSpecV1), ("v2",
+                                                     device_spec.DeviceSpecV2))
+
+
+class DeviceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(*TEST_V1_AND_V2)
+  def testMerge(self, DeviceSpec):  # pylint: disable=invalid-name
+    d = DeviceSpec.from_string("/job:muu/task:1/device:MyFunnyDevice:2")
+    self.assertEqual("/job:muu/task:1/device:MyFunnyDevice:2", d.to_string())
+
+    if not context.executing_eagerly():
+      with ops.device(device.merge_device("/device:GPU:0")):
+        var1 = variables.Variable(1.0)
+        self.assertEqual("/device:GPU:0", var1.device)
+        with ops.device(device.merge_device("/job:worker")):
+          var2 = variables.Variable(1.0)
+          self.assertEqual("/job:worker/device:GPU:0", var2.device)
+          with ops.device(device.merge_device("/device:CPU:0")):
+            var3 = variables.Variable(1.0)
+            self.assertEqual("/job:worker/device:CPU:0", var3.device)
+            with ops.device(device.merge_device("/job:ps")):
+              var4 = variables.Variable(1.0)
+              self.assertEqual("/job:ps/device:CPU:0", var4.device)
+
+  def testCanonicalName(self):
+    self.assertEqual("/job:foo/replica:0",
+                     device.canonical_name("/job:foo/replica:0"))
+    self.assertEqual("/job:foo/replica:0",
+                     device.canonical_name("/replica:0/job:foo"))
+
+    self.assertEqual("/job:foo/replica:0/task:0",
+                     device.canonical_name("/job:foo/replica:0/task:0"))
+    self.assertEqual("/job:foo/replica:0/task:0",
+                     device.canonical_name("/job:foo/task:0/replica:0"))
+
+    self.assertEqual("/device:CPU:0", device.canonical_name("/device:CPU:0"))
+    self.assertEqual("/device:GPU:2", device.canonical_name("/device:GPU:2"))
+
+    self.assertEqual(
+        "/job:foo/replica:0/task:0/device:GPU:0",
+        device.canonical_name("/job:foo/replica:0/task:0/device:GPU:0"))
+    self.assertEqual(
+        "/job:foo/replica:0/task:0/device:GPU:0",
+        device.canonical_name("/device:GPU:0/task:0/replica:0/job:foo"))
+
+  def testCheckValid(self):
+    device.check_valid("/job:foo/replica:0")
+
+    with self.assertRaisesRegex(ValueError, "invalid literal for int"):
+      device.check_valid("/job:j/replica:foo")
+
+    with self.assertRaisesRegex(ValueError, "invalid literal for int"):
+      device.check_valid("/job:j/task:bar")
+
+    # Assume no one will register a device type named "barcpugpu"
+    with self.assertRaisesRegex(ValueError, "Unknown attribute 'barcpugpu'"):
+      device.check_valid("/barcpugpu:muu/baz:2")
+
+    with self.assertRaisesRegex(ValueError,
+                                "Multiple device types are not allowed"):
+      device.check_valid("/cpu:0/device:GPU:2")
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/framework/dtypes.cc

@@ -0,0 +1,167 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "pybind11/detail/common.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace {
+
+inline int DataTypeId(tensorflow::DataType dt) { return static_cast<int>(dt); }
+
+// A variant of tensorflow::DataTypeString which uses fixed-width names
+// for floating point data types. This behavior is compatible with that of
+// existing pure Python DType.
+const std::string DataTypeStringCompat(tensorflow::DataType dt) {
+  switch (dt) {
+    case tensorflow::DataType::DT_HALF:
+      return "float16";
+    case tensorflow::DataType::DT_HALF_REF:
+      return "float16_ref";
+    case tensorflow::DataType::DT_FLOAT:
+      return "float32";
+    case tensorflow::DataType::DT_FLOAT_REF:
+      return "float32_ref";
+    case tensorflow::DataType::DT_DOUBLE:
+      return "float64";
+    case tensorflow::DataType::DT_DOUBLE_REF:
+      return "float64_ref";
+    default:
+      return tensorflow::DataTypeString(dt);
+  }
+}
+
+}  // namespace
+
+namespace tensorflow {
+
+constexpr DataTypeSet kNumPyIncompatibleTypes =
+    ToSet(DataType::DT_RESOURCE) | ToSet(DataType::DT_VARIANT);
+
+inline bool DataTypeIsNumPyCompatible(DataType dt) {
+  return !kNumPyIncompatibleTypes.Contains(dt);
+}
+
+}  // namespace tensorflow
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_dtypes, m, pybind11::mod_gil_not_used()) {
+  py::class_<tensorflow::DataType>(m, "DType")
+      .def(py::init([](py::object obj) {
+        auto id = static_cast<int>(py::int_(obj));
+        if (tensorflow::DataType_IsValid(id) &&
+            id != static_cast<int>(tensorflow::DT_INVALID)) {
+          return static_cast<tensorflow::DataType>(id);
+        }
+        throw py::type_error(
+            py::str("{} does not correspond to a valid tensorflow::DataType")
+                .format(id));
+      }))
+      .def("__int__",
+           [](tensorflow::DataType self) { return DataTypeId(self); })
+      // For compatibility with pure-Python DType.
+      .def_property_readonly("_type_enum", &DataTypeId)
+      .def_property_readonly(
+          "as_datatype_enum", &DataTypeId,
+          "Returns a `types_pb2.DataType` enum value based on this data type.")
+
+      .def_property_readonly("name",
+                             [](tensorflow::DataType self) {
+#if PY_MAJOR_VERSION < 3
+                               return py::bytes(DataTypeStringCompat(self));
+#else
+                               return DataTypeStringCompat(self);
+#endif
+                             })
+      .def_property_readonly(
+          "size",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeSize(tensorflow::BaseType(self));
+          })
+
+      .def("__repr__",
+           [](tensorflow::DataType self) {
+             return py::str("tf.{}").format(DataTypeStringCompat(self));
+           })
+      .def("__str__",
+           [](tensorflow::DataType self) {
+             return py::str("<dtype: {!r}>")
+#if PY_MAJOR_VERSION < 3
+                 .format(py::bytes(DataTypeStringCompat(self)));
+#else
+                 .format(DataTypeStringCompat(self));
+#endif
+           })
+      .def("__hash__", &DataTypeId)
+
+      .def_property_readonly(
+          "is_numpy_compatible",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsNumPyCompatible(
+                tensorflow::BaseType(self));
+          },
+          "Returns whether this data type has a compatible NumPy data type.")
+
+      .def_property_readonly(
+          "is_bool",
+          [](tensorflow::DataType self) {
+            return tensorflow::BaseType(self) == tensorflow::DT_BOOL;
+          },
+          "Returns whether this is a boolean data type.")
+      .def_property_readonly(
+          "is_numeric",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsNumeric(tensorflow::BaseType(self));
+          },
+          "Returns whether this is a numeric data type.")
+      .def_property_readonly(
+          "is_complex",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsComplex(tensorflow::BaseType(self));
+          },
+          "Returns whether this is a complex floating point type.")
+      .def_property_readonly(
+          "is_floating",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsFloating(tensorflow::BaseType(self));
+          },
+          "Returns whether this is a (non-quantized, real) floating point "
+          "type.")
+      .def_property_readonly(
+          "is_integer",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsInteger(tensorflow::BaseType(self));
+          },
+          "Returns whether this is a (non-quantized) integer type.")
+      .def_property_readonly(
+          "is_quantized",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsQuantized(tensorflow::BaseType(self));
+          },
+          "Returns whether this is a quantized data type.")
+      .def_property_readonly(
+          "is_unsigned",
+          [](tensorflow::DataType self) {
+            return tensorflow::DataTypeIsUnsigned(tensorflow::BaseType(self));
+          },
+          R"doc(Returns whether this type is unsigned.
+
+Non-numeric, unordered, and quantized types are not considered unsigned, and
+this function returns `False`.)doc");
+
+  py::implicitly_convertible<py::int_, tensorflow::DataType>();
+}

tensorflow/python/framework/dtypes.py

@@ -0,0 +1,945 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Library of dtypes (Tensor element types)."""
+import abc
+import builtins
+import dataclasses
+from typing import Type, Sequence, Optional
+
+import ml_dtypes
+import numpy as np
+
+from tensorflow.core.framework import types_pb2
+# We need to import pywrap_tensorflow prior to the bfloat wrapper to avoid
+# protobuf errors where a file is defined twice on MacOS.
+# pylint: disable=invalid-import-order,g-bad-import-order
+from tensorflow.python import pywrap_tensorflow  # pylint: disable=unused-import
+from tensorflow.python.framework import _dtypes
+from tensorflow.python.framework import cpp_shape_inference_pb2
+from tensorflow.python.types import doc_typealias
+from tensorflow.python.util.tf_export import tf_export
+from tensorflow.python.types import trace
+from tensorflow.core.function import trace_type
+from tensorflow.tools.docs import doc_controls
+
+
+class DTypeMeta(type(_dtypes.DType), abc.ABCMeta):
+  pass
+
+
+@dataclasses.dataclass(frozen=True)
+class HandleData:
+  """Holds resource/variant tensor specific data."""
+  shape_inference: Optional[
+      cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData
+  ] = None
+  alias_id: Optional[int] = None
+
+
+@tf_export("dtypes.DType", "DType")
+class DType(
+    _dtypes.DType,
+    trace.TraceType,
+    trace_type.Serializable,
+    metaclass=DTypeMeta):
+  """Represents the type of the elements in a `Tensor`.
+
+  `DType`'s are used to specify the output data type for operations which
+  require it, or to inspect the data type of existing `Tensor`'s.
+
+  Examples:
+
+  >>> tf.constant(1, dtype=tf.int64)
+  <tf.Tensor: shape=(), dtype=int64, numpy=1>
+  >>> tf.constant(1.0).dtype
+  tf.float32
+
+  See `tf.dtypes` for a complete list of `DType`'s defined.
+  """
+  __slots__ = ["_handle_data"]
+
+  def __init__(self, type_enum, handle_data=None):
+    super().__init__(type_enum)
+
+    # Resource and Variant dtypes have additional handle data information that
+    # is necessary for manipulating those Tensors.
+    if handle_data is not None and not isinstance(handle_data, HandleData):
+      raise TypeError("handle_data must be of the type HandleData")
+
+    self._handle_data = handle_data
+
+  @property
+  def _is_ref_dtype(self):
+    """Returns `True` if this `DType` represents a reference type."""
+    return self._type_enum > 100
+
+  @property
+  def _as_ref(self):
+    """Returns a reference `DType` based on this `DType`."""
+    if self._is_ref_dtype:
+      return self
+    else:
+      return _INTERN_TABLE[self._type_enum + 100]
+
+  @property
+  def base_dtype(self):
+    """Returns a non-reference `DType` based on this `DType` (for TF1).
+
+    Programs written for TensorFlow 2.x do not need this attribute.
+    It exists only for compatibility with TensorFlow 1.x, which used
+    reference `DType`s in the implementation of `tf.compat.v1.Variable`.
+    In TensorFlow 2.x, `tf.Variable` is implemented without reference types.
+    """
+    if self._is_ref_dtype:
+      return _INTERN_TABLE[self._type_enum - 100]
+    else:
+      return self
+
+  @property
+  def real_dtype(self):
+    """Returns the `DType` corresponding to this `DType`'s real part."""
+    base = self.base_dtype
+    if base == complex64:
+      return float32
+    elif base == complex128:
+      return float64
+    else:
+      return self
+
+  @property
+  def as_numpy_dtype(self):
+    """Returns a Python `type` object based on this `DType`."""
+    return _TF_TO_NP[self._type_enum]
+
+  @property
+  def min(self):
+    """Returns the minimum representable value in this data type.
+
+    Raises:
+      TypeError: if this is a non-numeric, unordered, or quantized type.
+
+    """
+    if (self.is_quantized or
+        self.base_dtype in (bool, string, complex64, complex128)):
+      raise TypeError(f"Cannot find minimum value of {self} with "
+                      f"{'quantized type' if self.is_quantized else 'type'} "
+                      f"{self.base_dtype}.")
+
+    # There is no simple way to get the min value of a dtype, we have to check
+    # float and int types separately.
+    try:
+      return ml_dtypes.finfo(self.as_numpy_dtype).min
+    except:  # bare except as possible raises by finfo not documented
+      try:
+        return ml_dtypes.iinfo(self.as_numpy_dtype).min
+      except:
+        raise TypeError(f"Cannot find minimum value of {self}.")
+
+  @property
+  def max(self):
+    """Returns the maximum representable value in this data type.
+
+    Raises:
+      TypeError: if this is a non-numeric, unordered, or quantized type.
+
+    """
+    if (self.is_quantized or
+        self.base_dtype in (bool, string, complex64, complex128)):
+      raise TypeError(f"Cannot find maximum value of {self} with "
+                      f"{'quantized type' if self.is_quantized else 'type'} "
+                      f"{self.base_dtype}.")
+
+    # there is no simple way to get the max value of a dtype, we have to check
+    # float and int types separately
+    try:
+      return ml_dtypes.finfo(self.as_numpy_dtype).max
+    except:  # bare except as possible raises by finfo not documented
+      try:
+        return ml_dtypes.iinfo(self.as_numpy_dtype).max
+      except:
+        raise TypeError(f"Cannot find maximum value of {self}.")
+
+  @property
+  def limits(self, clip_negative=True):
+    """Return intensity limits, i.e.
+
+    (min, max) tuple, of the dtype.
+    Args:
+      clip_negative : bool, optional If True, clip the negative range (i.e.
+        return 0 for min intensity) even if the image dtype allows negative
+        values. Returns
+      min, max : tuple Lower and upper intensity limits.
+    """
+    if self.as_numpy_dtype in dtype_range:
+      min, max = dtype_range[self.as_numpy_dtype]  # pylint: disable=redefined-builtin
+    else:
+      raise ValueError(str(self) + " does not have defined limits.")
+
+    if clip_negative:
+      min = 0  # pylint: disable=redefined-builtin
+    return min, max
+
+  def is_compatible_with(self, other):
+    """Returns True if the `other` DType will be converted to this DType (TF1).
+
+    Programs written for TensorFlow 2.x do not need this function.
+    Instead, they can do equality comparison on `DType` objects directly:
+    `tf.as_dtype(this) == tf.as_dtype(other)`.
+
+    This function exists only for compatibility with TensorFlow 1.x, where it
+    additionally allows conversion from a reference type (used by
+    `tf.compat.v1.Variable`) to its base type.
+
+    Args:
+      other: A `DType` (or object that may be converted to a `DType`).
+
+    Returns:
+      True if a Tensor of the `other` `DType` will be implicitly converted to
+      this `DType`.
+    """
+    other = as_dtype(other)
+    return self._type_enum in (other.as_datatype_enum,
+                               other.base_dtype.as_datatype_enum)
+
+  def is_subtype_of(self, other: trace.TraceType) -> bool:
+    """See tf.types.experimental.TraceType base class."""
+    return self == other
+
+  def most_specific_common_supertype(
+      self, types: Sequence[trace.TraceType]) -> Optional["DType"]:
+    """See tf.types.experimental.TraceType base class."""
+    return self if all(self == other for other in types) else None
+
+  @doc_controls.do_not_doc_inheritable
+  def placeholder_value(self, placeholder_context):
+    """See tf.types.experimental.TraceType base class."""
+    return super().placeholder_value(placeholder_context)
+
+  @doc_controls.do_not_doc_inheritable
+  def from_tensors(self, tensors):
+    """See tf.types.experimental.TraceType base class."""
+    return super().from_tensors(tensors)
+
+  @doc_controls.do_not_doc_inheritable
+  def to_tensors(self, value):
+    """See tf.types.experimental.TraceType base class."""
+    return super().to_tensors(value)
+
+  @doc_controls.do_not_doc_inheritable
+  def flatten(self):
+    """See tf.types.experimental.TraceType base class."""
+    return super().flatten()
+
+  @doc_controls.do_not_doc_inheritable
+  def cast(self, value, cast_context):
+    """See tf.types.experimental.TraceType base class."""
+    return super().cast(value, cast_context)
+
+  @classmethod
+  def experimental_type_proto(cls) -> Type[types_pb2.SerializedDType]:
+    """Returns the type of proto associated with DType serialization."""
+    return types_pb2.SerializedDType
+
+  @classmethod
+  def experimental_from_proto(cls, proto: types_pb2.SerializedDType) -> "DType":
+    """Returns a Dtype instance based on the serialized proto."""
+    return DType(proto.datatype)
+
+  def experimental_as_proto(self) -> types_pb2.SerializedDType:
+    """Returns a proto representation of the Dtype instance."""
+    return types_pb2.SerializedDType(datatype=self._type_enum)
+
+  def __eq__(self, other):
+    """Returns True iff this DType refers to the same type as `other`."""
+    if other is None:
+      return False
+
+    if type(other) != DType:  # pylint: disable=unidiomatic-typecheck
+      try:
+        other = as_dtype(other)
+      except TypeError:
+        return False
+
+    return self._type_enum == other._type_enum  # pylint: disable=protected-access
+
+  def __ne__(self, other):
+    """Returns True iff self != other."""
+    return not self.__eq__(other)
+
+  # "If a class that overrides __eq__() needs to retain the implementation
+  #  of __hash__() from a parent class, the interpreter must be told this
+  #  explicitly by setting __hash__ = <ParentClass>.__hash__."
+  # TODO(slebedev): Remove once __eq__ and __ne__ are implemented in C++.
+  __hash__ = _dtypes.DType.__hash__
+
+  def __reduce__(self):
+    return as_dtype, (self.name,)
+
+trace_type.register_serializable(DType)
+
+# Define data type range of numpy dtype
+dtype_range = {
+    np.bool_: (False, True),
+    np.uint8: (0, 255),
+    np.uint16: (0, 65535),
+    np.int8: (-128, 127),
+    np.int16: (-32768, 32767),
+    np.int64: (-(2**63), 2**63 - 1),
+    np.uint64: (0, 2**64 - 1),
+    np.int32: (-(2**31), 2**31 - 1),
+    np.uint32: (0, 2**32 - 1),
+    np.float32: (-1, 1),
+    np.float64: (-1, 1),
+}
+
+# Define standard wrappers for the types_pb2.DataType enum.
+resource = DType(types_pb2.DT_RESOURCE)
+doc_typealias.document(
+    obj=resource, doc="Handle to a mutable, dynamically allocated resource."
+)
+tf_export("dtypes.resource", "resource").export_constant(__name__, "resource")
+
+variant = DType(types_pb2.DT_VARIANT)
+doc_typealias.document(
+    obj=variant, doc="Data of arbitrary type (known at runtime)."
+)
+tf_export("dtypes.variant", "variant").export_constant(__name__, "variant")
+
+uint8 = DType(types_pb2.DT_UINT8)
+doc_typealias.document(obj=uint8, doc="Unsigned 8-bit (byte) integer.")
+tf_export("dtypes.uint8", "uint8").export_constant(__name__, "uint8")
+
+uint16 = DType(types_pb2.DT_UINT16)
+doc_typealias.document(obj=uint16, doc="Unsigned 16-bit (word) integer.")
+tf_export("dtypes.uint16", "uint16").export_constant(__name__, "uint16")
+
+uint32 = DType(types_pb2.DT_UINT32)
+doc_typealias.document(obj=uint32, doc="Unsigned 32-bit (dword) integer.")
+tf_export("dtypes.uint32", "uint32").export_constant(__name__, "uint32")
+
+uint64 = DType(types_pb2.DT_UINT64)
+doc_typealias.document(obj=uint64, doc="Unsigned 64-bit (qword) integer.")
+tf_export("dtypes.uint64", "uint64").export_constant(__name__, "uint64")
+
+int8 = DType(types_pb2.DT_INT8)
+doc_typealias.document(obj=int8, doc="Signed 8-bit integer.")
+tf_export("dtypes.int8", "int8").export_constant(__name__, "int8")
+
+int16 = DType(types_pb2.DT_INT16)
+doc_typealias.document(obj=int16, doc="Signed 16-bit integer.")
+tf_export("dtypes.int16", "int16").export_constant(__name__, "int16")
+
+int32 = DType(types_pb2.DT_INT32)
+doc_typealias.document(obj=int32, doc="Signed 32-bit integer.")
+tf_export("dtypes.int32", "int32").export_constant(__name__, "int32")
+
+int64 = DType(types_pb2.DT_INT64)
+doc_typealias.document(obj=int64, doc="Signed 64-bit integer.")
+tf_export("dtypes.int64", "int64").export_constant(__name__, "int64")
+
+float16 = DType(types_pb2.DT_HALF)
+half = float16
+doc_typealias.document(
+    obj=float16, doc="16-bit (half precision) floating-point."
+)
+tf_export("dtypes.float16", "float16").export_constant(__name__, "float16")
+tf_export("dtypes.half", "half").export_constant(__name__, "half")
+
+float32 = DType(types_pb2.DT_FLOAT)
+doc_typealias.document(
+    obj=float32, doc="32-bit (single precision) floating-point."
+)
+tf_export("dtypes.float32", "float32").export_constant(__name__, "float32")
+
+float64 = DType(types_pb2.DT_DOUBLE)
+doc_typealias.document(
+    obj=float64, doc="64-bit (double precision) floating-point."
+)
+tf_export("dtypes.float64", "float64").export_constant(__name__, "float64")
+double = float64
+tf_export("dtypes.double", "double").export_constant(__name__, "double")
+
+complex64 = DType(types_pb2.DT_COMPLEX64)
+doc_typealias.document(obj=complex64, doc="64-bit complex.")
+tf_export("dtypes.complex64", "complex64").export_constant(
+    __name__, "complex64"
+)
+
+complex128 = DType(types_pb2.DT_COMPLEX128)
+doc_typealias.document(obj=complex128, doc="128-bit complex.")
+tf_export("dtypes.complex128", "complex128").export_constant(
+    __name__, "complex128"
+)
+
+string = DType(types_pb2.DT_STRING)
+doc_typealias.document(
+    obj=string, doc="Variable-length string, represented as byte array."
+)
+tf_export("dtypes.string", "string").export_constant(__name__, "string")
+
+bool = DType(types_pb2.DT_BOOL)  # pylint: disable=redefined-builtin
+doc_typealias.document(obj=bool, doc="Boolean.")
+tf_export("dtypes.bool", "bool").export_constant(__name__, "bool")
+
+qint8 = DType(types_pb2.DT_QINT8)
+doc_typealias.document(obj=qint8, doc="Signed quantized 8-bit integer.")
+tf_export("dtypes.qint8", "qint8").export_constant(__name__, "qint8")
+
+qint16 = DType(types_pb2.DT_QINT16)
+doc_typealias.document(obj=qint16, doc="Signed quantized 16-bit integer.")
+tf_export("dtypes.qint16", "qint16").export_constant(__name__, "qint16")
+
+qint32 = DType(types_pb2.DT_QINT32)
+doc_typealias.document(obj=qint32, doc="signed quantized 32-bit integer.")
+tf_export("dtypes.qint32", "qint32").export_constant(__name__, "qint32")
+
+quint8 = DType(types_pb2.DT_QUINT8)
+doc_typealias.document(obj=quint8, doc="Unsigned quantized 8-bit integer.")
+tf_export("dtypes.quint8", "quint8").export_constant(__name__, "quint8")
+
+quint16 = DType(types_pb2.DT_QUINT16)
+doc_typealias.document(obj=quint16, doc="Unsigned quantized 16-bit integer.")
+tf_export("dtypes.quint16", "quint16").export_constant(__name__, "quint16")
+
+bfloat16 = DType(types_pb2.DT_BFLOAT16)
+doc_typealias.document(
+    obj=bfloat16, doc="16-bit bfloat (brain floating point)."
+)
+tf_export("dtypes.bfloat16", "bfloat16").export_constant(__name__, "bfloat16")
+
+float8_e5m2 = DType(types_pb2.DT_FLOAT8_E5M2)
+doc_typealias.document(
+    obj=float8_e5m2, doc="8-bit float with 5 exponent bits and 2 mantissa bits."
+)
+tf_export(
+    "dtypes.experimental.float8_e5m2", "experimental.float8_e5m2"
+).export_constant(__name__, "float8_e5m2")
+
+float8_e4m3fn = DType(types_pb2.DT_FLOAT8_E4M3FN)
+doc_typealias.document(
+    obj=float8_e4m3fn,
+    doc=(
+        "8-bit float with 4 exponent bits and 3 mantissa bits, with extended"
+        " finite range.  This type has no representation for inf, and only two"
+        " NaN values: 0xFF for negative NaN, and 0x7F for positive NaN."
+    ),
+)
+tf_export(
+    "dtypes.experimental.float8_e4m3fn", "experimental.float8_e4m3fn"
+).export_constant(__name__, "float8_e4m3fn")
+
+float8_e4m3fnuz = DType(types_pb2.DT_FLOAT8_E4M3FNUZ)
+doc_typealias.document(
+    obj=float8_e4m3fnuz,
+    doc=(
+        "8-bit float with 4 exponent bits and 3 mantissa bits, with extended"
+        " finite range.  This type has no representation for inf, and only one"
+        " NaN value: 0x80."
+    ),
+)
+tf_export(
+    "dtypes.experimental.float8_e4m3fnuz", "experimental.float8_e4m3fnuz"
+).export_constant(__name__, "float8_e4m3fnuz")
+
+float8_e4m3b11fnuz = DType(types_pb2.DT_FLOAT8_E4M3B11FNUZ)
+doc_typealias.document(
+    obj=float8_e4m3b11fnuz,
+    doc=(
+        "8-bit float with 4 exponent bits and 3 mantissa bits, with extended "
+        "finite range and 11 bits of bias.  This type has no representation "
+        "for inf, and only one NaN value: 0x80."
+    ),
+)
+tf_export(
+    "dtypes.experimental.float8_e4m3b11fnuz", "experimental.float8_e4m3b11fnuz"
+).export_constant(__name__, "float8_e4m3b11fnuz")
+
+float8_e5m2fnuz = DType(types_pb2.DT_FLOAT8_E5M2FNUZ)
+doc_typealias.document(
+    obj=float8_e5m2fnuz,
+    doc=(
+        "8-bit float with 5 exponent bits and 2 mantissa bits, with extended "
+        "finite range.  This type has no representation for inf, and only one "
+        "NaN value: 0x80."
+    ),
+)
+tf_export(
+    "dtypes.experimental.float8_e5m2fnuz", "experimental.float8_e5m2fnuz"
+).export_constant(__name__, "float8_e5m2fnuz")
+
+float4_e2m1fn = DType(types_pb2.DT_FLOAT4_E2M1FN)
+doc_typealias.document(
+    obj=float4_e2m1fn,
+    doc=(
+        "4-bit float with 2 exponent bits and 1 mantissa bits, with extended"
+        " finite range.  This type has no representation for both inf and NaN."
+    ),
+)
+tf_export(
+    "dtypes.experimental.float4_e2m1fn", "experimental.float4_e2m1fn"
+).export_constant(__name__, "float4_e2m1fn")
+
+
+int4 = DType(types_pb2.DT_INT4)
+doc_typealias.document(obj=int4, doc="Signed 4-bit integer.")
+tf_export("dtypes.experimental.int4", "experimental.int4").export_constant(
+    __name__, "int4"
+)
+
+uint4 = DType(types_pb2.DT_UINT4)
+doc_typealias.document(obj=uint4, doc="Unsigned 4-bit integer.")
+tf_export("dtypes.experimental.uint4", "experimental.uint4").export_constant(
+    __name__, "uint4"
+)
+
+int2 = DType(types_pb2.DT_INT2)
+doc_typealias.document(obj=int2, doc="Signed 2-bit integer.")
+tf_export("dtypes.experimental.int2", "experimental.int2").export_constant(
+    __name__, "int2"
+)
+
+uint2 = DType(types_pb2.DT_UINT2)
+doc_typealias.document(obj=uint2, doc="Unsigned 2-bit integer.")
+tf_export("dtypes.experimental.uint2", "experimental.uint2").export_constant(
+    __name__, "uint2"
+)
+
+resource_ref = DType(types_pb2.DT_RESOURCE_REF)
+variant_ref = DType(types_pb2.DT_VARIANT_REF)
+float16_ref = DType(types_pb2.DT_HALF_REF)
+half_ref = float16_ref
+float32_ref = DType(types_pb2.DT_FLOAT_REF)
+float64_ref = DType(types_pb2.DT_DOUBLE_REF)
+double_ref = float64_ref
+int32_ref = DType(types_pb2.DT_INT32_REF)
+uint32_ref = DType(types_pb2.DT_UINT32_REF)
+uint8_ref = DType(types_pb2.DT_UINT8_REF)
+uint16_ref = DType(types_pb2.DT_UINT16_REF)
+int16_ref = DType(types_pb2.DT_INT16_REF)
+int8_ref = DType(types_pb2.DT_INT8_REF)
+string_ref = DType(types_pb2.DT_STRING_REF)
+complex64_ref = DType(types_pb2.DT_COMPLEX64_REF)
+complex128_ref = DType(types_pb2.DT_COMPLEX128_REF)
+int64_ref = DType(types_pb2.DT_INT64_REF)
+uint64_ref = DType(types_pb2.DT_UINT64_REF)
+bool_ref = DType(types_pb2.DT_BOOL_REF)
+qint8_ref = DType(types_pb2.DT_QINT8_REF)
+quint8_ref = DType(types_pb2.DT_QUINT8_REF)
+qint16_ref = DType(types_pb2.DT_QINT16_REF)
+quint16_ref = DType(types_pb2.DT_QUINT16_REF)
+qint32_ref = DType(types_pb2.DT_QINT32_REF)
+bfloat16_ref = DType(types_pb2.DT_BFLOAT16_REF)
+float8_e5m2_ref = DType(types_pb2.DT_FLOAT8_E5M2_REF)
+float8_e4m3fn_ref = DType(types_pb2.DT_FLOAT8_E4M3FN_REF)
+float8_e4m3fnuz_ref = DType(types_pb2.DT_FLOAT8_E4M3FNUZ_REF)
+float8_e4m3b11fnuz_ref = DType(types_pb2.DT_FLOAT8_E4M3B11FNUZ_REF)
+float8_e5m2fnuz_ref = DType(types_pb2.DT_FLOAT8_E5M2FNUZ_REF)
+float4_e2m1fn_ref = DType(types_pb2.DT_FLOAT4_E2M1FN_REF)
+int4_ref = DType(types_pb2.DT_INT4_REF)
+uint4_ref = DType(types_pb2.DT_UINT4_REF)
+int2_ref = DType(types_pb2.DT_INT2_REF)
+uint2_ref = DType(types_pb2.DT_UINT2_REF)
+
+# Maintain an intern table so that we don't have to create a large
+# number of small objects.
+_INTERN_TABLE = {
+    types_pb2.DT_HALF: float16,
+    types_pb2.DT_FLOAT: float32,
+    types_pb2.DT_DOUBLE: float64,
+    types_pb2.DT_INT32: int32,
+    types_pb2.DT_UINT8: uint8,
+    types_pb2.DT_UINT16: uint16,
+    types_pb2.DT_UINT32: uint32,
+    types_pb2.DT_UINT64: uint64,
+    types_pb2.DT_INT16: int16,
+    types_pb2.DT_INT8: int8,
+    types_pb2.DT_STRING: string,
+    types_pb2.DT_COMPLEX64: complex64,
+    types_pb2.DT_COMPLEX128: complex128,
+    types_pb2.DT_INT64: int64,
+    types_pb2.DT_BOOL: bool,
+    types_pb2.DT_QINT8: qint8,
+    types_pb2.DT_QUINT8: quint8,
+    types_pb2.DT_QINT16: qint16,
+    types_pb2.DT_QUINT16: quint16,
+    types_pb2.DT_QINT32: qint32,
+    types_pb2.DT_BFLOAT16: bfloat16,
+    types_pb2.DT_FLOAT8_E5M2: float8_e5m2,
+    types_pb2.DT_FLOAT8_E4M3FN: float8_e4m3fn,
+    types_pb2.DT_FLOAT8_E4M3FNUZ: float8_e4m3fnuz,
+    types_pb2.DT_FLOAT8_E4M3B11FNUZ: float8_e4m3b11fnuz,
+    types_pb2.DT_FLOAT8_E5M2FNUZ: float8_e5m2fnuz,
+    types_pb2.DT_FLOAT4_E2M1FN: float4_e2m1fn,
+    types_pb2.DT_INT4: int4,
+    types_pb2.DT_UINT4: uint4,
+    types_pb2.DT_INT2: int2,
+    types_pb2.DT_UINT2: uint2,
+    types_pb2.DT_RESOURCE: resource,
+    types_pb2.DT_VARIANT: variant,
+    types_pb2.DT_HALF_REF: float16_ref,
+    types_pb2.DT_FLOAT_REF: float32_ref,
+    types_pb2.DT_DOUBLE_REF: float64_ref,
+    types_pb2.DT_INT32_REF: int32_ref,
+    types_pb2.DT_UINT32_REF: uint32_ref,
+    types_pb2.DT_UINT8_REF: uint8_ref,
+    types_pb2.DT_UINT16_REF: uint16_ref,
+    types_pb2.DT_INT16_REF: int16_ref,
+    types_pb2.DT_INT8_REF: int8_ref,
+    types_pb2.DT_STRING_REF: string_ref,
+    types_pb2.DT_COMPLEX64_REF: complex64_ref,
+    types_pb2.DT_COMPLEX128_REF: complex128_ref,
+    types_pb2.DT_INT64_REF: int64_ref,
+    types_pb2.DT_UINT64_REF: uint64_ref,
+    types_pb2.DT_BOOL_REF: bool_ref,
+    types_pb2.DT_QINT8_REF: qint8_ref,
+    types_pb2.DT_QUINT8_REF: quint8_ref,
+    types_pb2.DT_QINT16_REF: qint16_ref,
+    types_pb2.DT_QUINT16_REF: quint16_ref,
+    types_pb2.DT_QINT32_REF: qint32_ref,
+    types_pb2.DT_BFLOAT16_REF: bfloat16_ref,
+    types_pb2.DT_FLOAT8_E5M2_REF: float8_e5m2_ref,
+    types_pb2.DT_FLOAT8_E4M3FN_REF: float8_e4m3fn_ref,
+    types_pb2.DT_FLOAT8_E4M3FNUZ_REF: float8_e4m3fnuz_ref,
+    types_pb2.DT_FLOAT8_E4M3B11FNUZ_REF: float8_e4m3b11fnuz_ref,
+    types_pb2.DT_FLOAT8_E5M2FNUZ_REF: float8_e5m2fnuz_ref,
+    types_pb2.DT_FLOAT4_E2M1FN_REF: float4_e2m1fn_ref,
+    types_pb2.DT_INT4_REF: int4_ref,
+    types_pb2.DT_UINT4_REF: uint4_ref,
+    types_pb2.DT_INT2_REF: int2_ref,
+    types_pb2.DT_UINT2_REF: uint2_ref,
+    types_pb2.DT_RESOURCE_REF: resource_ref,
+    types_pb2.DT_VARIANT_REF: variant_ref,
+}
+
+# Standard mappings between types_pb2.DataType values and string names.
+_TYPE_TO_STRING = {
+    types_pb2.DT_HALF: "float16",
+    types_pb2.DT_FLOAT: "float32",
+    types_pb2.DT_DOUBLE: "float64",
+    types_pb2.DT_INT32: "int32",
+    types_pb2.DT_UINT8: "uint8",
+    types_pb2.DT_UINT16: "uint16",
+    types_pb2.DT_UINT32: "uint32",
+    types_pb2.DT_UINT64: "uint64",
+    types_pb2.DT_INT16: "int16",
+    types_pb2.DT_INT8: "int8",
+    types_pb2.DT_STRING: "string",
+    types_pb2.DT_COMPLEX64: "complex64",
+    types_pb2.DT_COMPLEX128: "complex128",
+    types_pb2.DT_INT64: "int64",
+    types_pb2.DT_BOOL: "bool",
+    types_pb2.DT_QINT8: "qint8",
+    types_pb2.DT_QUINT8: "quint8",
+    types_pb2.DT_QINT16: "qint16",
+    types_pb2.DT_QUINT16: "quint16",
+    types_pb2.DT_QINT32: "qint32",
+    types_pb2.DT_BFLOAT16: "bfloat16",
+    types_pb2.DT_FLOAT8_E5M2: "float8_e5m2",
+    types_pb2.DT_FLOAT8_E4M3FN: "float8_e4m3fn",
+    types_pb2.DT_FLOAT8_E4M3FNUZ: "float8_e4m3fnuz",
+    types_pb2.DT_FLOAT8_E4M3B11FNUZ: "float8_e4m3b11fnuz",
+    types_pb2.DT_FLOAT8_E5M2FNUZ: "float8_e5m2fnuz",
+    types_pb2.DT_FLOAT4_E2M1FN: "float4_e2m1fn",
+    types_pb2.DT_INT4: "int4",
+    types_pb2.DT_UINT4: "uint4",
+    types_pb2.DT_INT2: "int2",
+    types_pb2.DT_UINT2: "uint2",
+    types_pb2.DT_RESOURCE: "resource",
+    types_pb2.DT_VARIANT: "variant",
+    types_pb2.DT_HALF_REF: "float16_ref",
+    types_pb2.DT_FLOAT_REF: "float32_ref",
+    types_pb2.DT_DOUBLE_REF: "float64_ref",
+    types_pb2.DT_INT32_REF: "int32_ref",
+    types_pb2.DT_UINT32_REF: "uint32_ref",
+    types_pb2.DT_UINT8_REF: "uint8_ref",
+    types_pb2.DT_UINT16_REF: "uint16_ref",
+    types_pb2.DT_INT16_REF: "int16_ref",
+    types_pb2.DT_INT8_REF: "int8_ref",
+    types_pb2.DT_STRING_REF: "string_ref",
+    types_pb2.DT_COMPLEX64_REF: "complex64_ref",
+    types_pb2.DT_COMPLEX128_REF: "complex128_ref",
+    types_pb2.DT_INT64_REF: "int64_ref",
+    types_pb2.DT_UINT64_REF: "uint64_ref",
+    types_pb2.DT_BOOL_REF: "bool_ref",
+    types_pb2.DT_QINT8_REF: "qint8_ref",
+    types_pb2.DT_QUINT8_REF: "quint8_ref",
+    types_pb2.DT_QINT16_REF: "qint16_ref",
+    types_pb2.DT_QUINT16_REF: "quint16_ref",
+    types_pb2.DT_QINT32_REF: "qint32_ref",
+    types_pb2.DT_BFLOAT16_REF: "bfloat16_ref",
+    types_pb2.DT_FLOAT8_E5M2_REF: "float8_e5m2_ref",
+    types_pb2.DT_FLOAT8_E4M3FN_REF: "float8_e4m3fn_ref",
+    types_pb2.DT_FLOAT8_E4M3FNUZ_REF: "float8_e4m3fnuz_ref",
+    types_pb2.DT_FLOAT8_E4M3B11FNUZ_REF: "float8_e4m3b11fnuz_ref",
+    types_pb2.DT_FLOAT8_E5M2FNUZ_REF: "float8_e5m2fnuz_ref",
+    types_pb2.DT_FLOAT4_E2M1FN_REF: "float4_e2m1fn_ref",
+    types_pb2.DT_INT4_REF: "int4_ref",
+    types_pb2.DT_UINT4_REF: "uint4_ref",
+    types_pb2.DT_INT2_REF: "int2_ref",
+    types_pb2.DT_UINT2_REF: "uint2_ref",
+    types_pb2.DT_RESOURCE_REF: "resource_ref",
+    types_pb2.DT_VARIANT_REF: "variant_ref",
+}
+_STRING_TO_TF = {
+    value: _INTERN_TABLE[key] for key, value in _TYPE_TO_STRING.items()
+}
+# Add non-canonical aliases.
+_STRING_TO_TF["half"] = float16
+_STRING_TO_TF["half_ref"] = float16_ref
+_STRING_TO_TF["float"] = float32
+_STRING_TO_TF["float_ref"] = float32_ref
+_STRING_TO_TF["double"] = float64
+_STRING_TO_TF["double_ref"] = float64_ref
+
+# Numpy representation for quantized dtypes.
+#
+# These are magic strings that are used in the swig wrapper to identify
+# quantized types.
+# TODO(mrry,keveman): Investigate Numpy type registration to replace this
+# hard-coding of names.
+_np_qint8 = np.dtype([("qint8", np.int8)])
+_np_quint8 = np.dtype([("quint8", np.uint8)])
+_np_qint16 = np.dtype([("qint16", np.int16)])
+_np_quint16 = np.dtype([("quint16", np.uint16)])
+_np_qint32 = np.dtype([("qint32", np.int32)])
+
+# _np_bfloat16, _np_float8*, _np_(u)int4 are defined by module imports.
+_np_bfloat16 = ml_dtypes.bfloat16
+_np_float8_e4m3fn = ml_dtypes.float8_e4m3fn
+_np_float8_e4m3fnuz = ml_dtypes.float8_e4m3fnuz
+_np_float8_e4m3b11fnuz = ml_dtypes.float8_e4m3b11fnuz
+_np_float8_e5m2 = ml_dtypes.float8_e5m2
+_np_float8_e5m2fnuz = ml_dtypes.float8_e5m2fnuz
+_np_float4_e2m1fn = ml_dtypes.float4_e2m1fn
+_np_int4 = ml_dtypes.int4
+_np_uint4 = ml_dtypes.uint4
+_np_int2 = ml_dtypes.int2
+_np_uint2 = ml_dtypes.uint2
+
+# Custom struct dtype for directly-fed ResourceHandles of supported type(s).
+np_resource = np.dtype([("resource", np.ubyte)])
+
+# Standard mappings between types_pb2.DataType values and numpy.dtypes.
+_NP_TO_TF = {
+    np.float16: float16,
+    np.float32: float32,
+    np.float64: float64,
+    np.int32: int32,
+    np.int64: int64,
+    np.uint8: uint8,
+    np.uint16: uint16,
+    np.uint32: uint32,
+    np.uint64: uint64,
+    np.int16: int16,
+    np.int8: int8,
+    np.complex64: complex64,
+    np.complex128: complex128,
+    np.object_: string,
+    np.bytes_: string,
+    np.str_: string,
+    np.bool_: bool,
+    _np_qint8: qint8,
+    _np_quint8: quint8,
+    _np_qint16: qint16,
+    _np_quint16: quint16,
+    _np_qint32: qint32,
+    _np_bfloat16: bfloat16,
+    _np_float8_e5m2: float8_e5m2,
+    _np_float8_e4m3fn: float8_e4m3fn,
+    _np_float8_e4m3fnuz: float8_e4m3fnuz,
+    _np_float8_e4m3b11fnuz: float8_e4m3b11fnuz,
+    _np_float8_e5m2fnuz: float8_e5m2fnuz,
+    _np_float4_e2m1fn: float4_e2m1fn,
+    _np_int4: int4,
+    _np_uint4: uint4,
+    _np_int2: int2,
+    _np_uint2: uint2,
+}
+
+# Map (some) NumPy platform dtypes to TF ones using their fixed-width
+# synonyms. Note that platform dtypes are not always simples aliases,
+# i.e. reference equality is not guaranteed. See e.g. numpy/numpy#9799.
+for pdt in [
+    np.intc,
+    np.uintc,
+    np.int_,
+    np.uint,
+    np.longlong,
+    np.ulonglong,
+]:
+  if pdt not in _NP_TO_TF:
+    _NP_TO_TF[pdt] = next(
+        _NP_TO_TF[dt] for dt in _NP_TO_TF if dt == pdt().dtype)  # pylint: disable=no-value-for-parameter
+
+if hasattr(np.dtypes, "StringDType"):
+  _NP_TO_TF[np.dtypes.StringDType()] = string
+
+TF_VALUE_DTYPES = set(_NP_TO_TF.values())
+
+_TF_TO_NP = {
+    types_pb2.DT_HALF: np.float16,
+    types_pb2.DT_FLOAT: np.float32,
+    types_pb2.DT_DOUBLE: np.float64,
+    types_pb2.DT_INT32: np.int32,
+    types_pb2.DT_UINT8: np.uint8,
+    types_pb2.DT_UINT16: np.uint16,
+    types_pb2.DT_UINT32: np.uint32,
+    types_pb2.DT_UINT64: np.uint64,
+    types_pb2.DT_INT16: np.int16,
+    types_pb2.DT_INT8: np.int8,
+    # NOTE(touts): For strings we use object as it supports variable length
+    # strings.
+    types_pb2.DT_STRING: object,
+    types_pb2.DT_COMPLEX64: np.complex64,
+    types_pb2.DT_COMPLEX128: np.complex128,
+    types_pb2.DT_INT64: np.int64,
+    types_pb2.DT_BOOL: np.bool_,
+    types_pb2.DT_QINT8: _np_qint8,
+    types_pb2.DT_QUINT8: _np_quint8,
+    types_pb2.DT_QINT16: _np_qint16,
+    types_pb2.DT_QUINT16: _np_quint16,
+    types_pb2.DT_QINT32: _np_qint32,
+    types_pb2.DT_BFLOAT16: _np_bfloat16,
+    types_pb2.DT_FLOAT8_E5M2: _np_float8_e5m2,
+    types_pb2.DT_FLOAT8_E4M3FN: _np_float8_e4m3fn,
+    types_pb2.DT_FLOAT8_E4M3FNUZ: _np_float8_e4m3fnuz,
+    types_pb2.DT_FLOAT8_E4M3B11FNUZ: _np_float8_e4m3b11fnuz,
+    types_pb2.DT_FLOAT8_E5M2FNUZ: _np_float8_e5m2fnuz,
+    types_pb2.DT_FLOAT4_E2M1FN: _np_float4_e2m1fn,
+    types_pb2.DT_INT4: _np_int4,
+    types_pb2.DT_UINT4: _np_uint4,
+    types_pb2.DT_INT2: _np_int2,
+    types_pb2.DT_UINT2: _np_uint2,
+    # Ref types
+    types_pb2.DT_HALF_REF: np.float16,
+    types_pb2.DT_FLOAT_REF: np.float32,
+    types_pb2.DT_DOUBLE_REF: np.float64,
+    types_pb2.DT_INT32_REF: np.int32,
+    types_pb2.DT_UINT32_REF: np.uint32,
+    types_pb2.DT_UINT8_REF: np.uint8,
+    types_pb2.DT_UINT16_REF: np.uint16,
+    types_pb2.DT_INT16_REF: np.int16,
+    types_pb2.DT_INT8_REF: np.int8,
+    types_pb2.DT_STRING_REF: np.object_,
+    types_pb2.DT_COMPLEX64_REF: np.complex64,
+    types_pb2.DT_COMPLEX128_REF: np.complex128,
+    types_pb2.DT_INT64_REF: np.int64,
+    types_pb2.DT_UINT64_REF: np.uint64,
+    types_pb2.DT_BOOL_REF: np.bool_,
+    types_pb2.DT_QINT8_REF: _np_qint8,
+    types_pb2.DT_QUINT8_REF: _np_quint8,
+    types_pb2.DT_QINT16_REF: _np_qint16,
+    types_pb2.DT_QUINT16_REF: _np_quint16,
+    types_pb2.DT_QINT32_REF: _np_qint32,
+    types_pb2.DT_BFLOAT16_REF: _np_bfloat16,
+    types_pb2.DT_FLOAT8_E5M2_REF: _np_float8_e5m2,
+    types_pb2.DT_FLOAT8_E4M3FN_REF: _np_float8_e4m3fn,
+    types_pb2.DT_FLOAT8_E4M3FNUZ_REF: _np_float8_e4m3fnuz,
+    types_pb2.DT_FLOAT8_E4M3B11FNUZ_REF: _np_float8_e4m3b11fnuz,
+    types_pb2.DT_FLOAT8_E5M2FNUZ_REF: _np_float8_e5m2fnuz,
+    types_pb2.DT_FLOAT4_E2M1FN_REF: _np_float4_e2m1fn,
+    types_pb2.DT_INT4_REF: _np_int4,
+    types_pb2.DT_UINT4_REF: _np_uint4,
+    types_pb2.DT_INT2_REF: _np_int2,
+    types_pb2.DT_UINT2_REF: _np_uint2,
+}
+
+_QUANTIZED_DTYPES_NO_REF = frozenset([qint8, quint8, qint16, quint16, qint32])
+_QUANTIZED_DTYPES_REF = frozenset(
+    [qint8_ref, quint8_ref, qint16_ref, quint16_ref, qint32_ref])
+QUANTIZED_DTYPES = _QUANTIZED_DTYPES_REF.union(_QUANTIZED_DTYPES_NO_REF)
+tf_export(
+    "dtypes.QUANTIZED_DTYPES",
+    v1=["dtypes.QUANTIZED_DTYPES",
+        "QUANTIZED_DTYPES"]).export_constant(__name__, "QUANTIZED_DTYPES")
+
+_PYTHON_TO_TF = {
+    builtins.float: float32,
+    builtins.bool: bool,
+    builtins.object: string
+}
+
+_ANY_TO_TF = {}
+_ANY_TO_TF.update(_INTERN_TABLE)
+_ANY_TO_TF.update(_STRING_TO_TF)
+_ANY_TO_TF.update(_PYTHON_TO_TF)
+_ANY_TO_TF.update(_NP_TO_TF)
+
+# Ensure no collisions.
+assert len(_ANY_TO_TF) == sum(
+    len(d) for d in [_INTERN_TABLE, _STRING_TO_TF, _PYTHON_TO_TF, _NP_TO_TF])
+
+
+@tf_export("dtypes.as_dtype", "as_dtype")
+def as_dtype(type_value):
+  """Converts the given `type_value` to a `tf.DType`.
+
+  Inputs can be existing `tf.DType` objects, a [`DataType`
+  enum](https://www.tensorflow.org/code/tensorflow/core/framework/types.proto),
+  a string type name, or a
+  [`numpy.dtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.html).
+
+  Examples:
+  >>> tf.as_dtype(2)  # Enum value for float64.
+  tf.float64
+
+  >>> tf.as_dtype('float')
+  tf.float32
+
+  >>> tf.as_dtype(np.int32)
+  tf.int32
+
+  Note: `DType` values are interned (i.e. a single instance of each dtype is
+  stored in a map). When passed a new `DType` object, `as_dtype` always returns
+  the interned value.
+
+  Args:
+    type_value: A value that can be converted to a `tf.DType` object.
+
+  Returns:
+    A `DType` corresponding to `type_value`.
+
+  Raises:
+    TypeError: If `type_value` cannot be converted to a `DType`.
+  """
+  if isinstance(type_value, DType):
+    if type_value._handle_data is None:  # pylint:disable=protected-access
+      return _INTERN_TABLE[type_value.as_datatype_enum]
+    else:
+      return type_value
+
+  if isinstance(type_value, np.dtype):
+    try:
+      return _NP_TO_TF[type_value.type]
+    except KeyError:
+      pass
+
+  try:
+    return _ANY_TO_TF[type_value]
+  except (KeyError, TypeError):
+    # TypeError indicates that type_value is not hashable.
+    pass
+
+  if hasattr(type_value, "dtype"):
+    try:
+      return _NP_TO_TF[np.dtype(type_value.dtype).type]
+    except (KeyError, TypeError):
+      pass
+
+  if isinstance(type_value, _dtypes.DType):
+    return _INTERN_TABLE[type_value.as_datatype_enum]
+
+  raise TypeError(f"Cannot convert the argument `type_value`: {type_value!r} "
+                  "to a TensorFlow DType.")

tensorflow/python/framework/dtypes_test.py

@@ -0,0 +1,494 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.dtypes."""
+
+from absl.testing import parameterized
+import numpy as np
+
+# pylint: disable=g-bad-import-order
+from tensorflow.python.framework import _dtypes
+# pylint: enable=g-bad-import-order
+
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.function import trace_type
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+
+def _is_numeric_dtype_enum(datatype_enum):
+  non_numeric_dtypes = [
+      types_pb2.DT_VARIANT, types_pb2.DT_VARIANT_REF, types_pb2.DT_INVALID,
+      types_pb2.DT_RESOURCE, types_pb2.DT_RESOURCE_REF
+  ]
+  return datatype_enum not in non_numeric_dtypes
+
+
+class TypesTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  def testAllTypesConstructible(self):
+    for datatype_enum in types_pb2.DataType.values():
+      if datatype_enum == types_pb2.DT_INVALID:
+        continue
+      self.assertEqual(datatype_enum,
+                       dtypes.DType(datatype_enum).as_datatype_enum)
+
+  def testAllTypesConvertibleToDType(self):
+    for datatype_enum in types_pb2.DataType.values():
+      if datatype_enum == types_pb2.DT_INVALID:
+        continue
+      dt = dtypes.as_dtype(datatype_enum)
+      self.assertEqual(datatype_enum, dt.as_datatype_enum)
+
+  def testAllTypesConvertibleToNumpyDtype(self):
+    for datatype_enum in types_pb2.DataType.values():
+      if not _is_numeric_dtype_enum(datatype_enum):
+        continue
+      dtype = dtypes.as_dtype(datatype_enum)
+      numpy_dtype = dtype.as_numpy_dtype
+      _ = np.empty((1, 1, 1, 1), dtype=numpy_dtype)
+      if dtype.base_dtype != dtypes.bfloat16:
+        # NOTE(touts): Intentionally no way to feed a DT_BFLOAT16.
+        self.assertEqual(
+            dtypes.as_dtype(datatype_enum).base_dtype,
+            dtypes.as_dtype(numpy_dtype))
+
+  def testAllPybind11DTypeConvertibleToDType(self):
+    for datatype_enum in types_pb2.DataType.values():
+      if datatype_enum == types_pb2.DT_INVALID:
+        continue
+      dtype = _dtypes.DType(datatype_enum)
+      self.assertEqual(dtypes.as_dtype(datatype_enum), dtype)
+
+  def testInvalid(self):
+    with self.assertRaises(TypeError):
+      dtypes.DType(types_pb2.DT_INVALID)
+    with self.assertRaises(TypeError):
+      dtypes.as_dtype(types_pb2.DT_INVALID)
+
+  def testNumpyConversion(self):
+    self.assertIs(dtypes.float32, dtypes.as_dtype(np.float32))
+    self.assertIs(dtypes.float64, dtypes.as_dtype(np.float64))
+    self.assertIs(dtypes.int32, dtypes.as_dtype(np.int32))
+    self.assertIs(dtypes.int64, dtypes.as_dtype(np.int64))
+    self.assertIs(dtypes.uint8, dtypes.as_dtype(np.uint8))
+    self.assertIs(dtypes.uint16, dtypes.as_dtype(np.uint16))
+    self.assertIs(dtypes.int16, dtypes.as_dtype(np.int16))
+    self.assertIs(dtypes.int8, dtypes.as_dtype(np.int8))
+    self.assertIs(dtypes.complex64, dtypes.as_dtype(np.complex64))
+    self.assertIs(dtypes.complex128, dtypes.as_dtype(np.complex128))
+    self.assertIs(dtypes.string, dtypes.as_dtype(np.object_))
+    self.assertIs(dtypes.string,
+                  dtypes.as_dtype(np.array(["foo", "bar"]).dtype))
+    self.assertIs(dtypes.bool, dtypes.as_dtype(np.bool_))
+    self.assertIs(dtypes.float8_e5m2, dtypes.as_dtype(dtypes._np_float8_e5m2))
+    self.assertIs(dtypes.float8_e4m3fn,
+                  dtypes.as_dtype(dtypes._np_float8_e4m3fn))
+    self.assertIs(
+        dtypes.float4_e2m1fn, dtypes.as_dtype(dtypes._np_float4_e2m1fn)
+    )
+    self.assertIs(
+        dtypes.float8_e4m3fnuz, dtypes.as_dtype(dtypes._np_float8_e4m3fnuz)
+    )
+    self.assertIs(
+        dtypes.float8_e4m3b11fnuz,
+        dtypes.as_dtype(dtypes._np_float8_e4m3b11fnuz),
+    )
+    self.assertIs(
+        dtypes.float8_e5m2fnuz, dtypes.as_dtype(dtypes._np_float8_e5m2fnuz)
+    )
+    self.assertIs(dtypes.int4, dtypes.as_dtype(dtypes._np_int4))
+    self.assertIs(dtypes.uint4, dtypes.as_dtype(dtypes._np_uint4))
+    self.assertIs(dtypes.int2, dtypes.as_dtype(dtypes._np_int2))
+    self.assertIs(dtypes.uint2, dtypes.as_dtype(dtypes._np_uint2))
+    with self.assertRaises(TypeError):
+      dtypes.as_dtype(np.dtype([("f1", np.uint), ("f2", np.int32)]))
+
+    class AnObject(object):
+      dtype = "f4"
+
+    self.assertIs(dtypes.float32, dtypes.as_dtype(AnObject))
+
+    class AnotherObject(object):
+      dtype = np.dtype(np.complex64)
+
+    self.assertIs(dtypes.complex64, dtypes.as_dtype(AnotherObject))
+
+  def testRealDtype(self):
+    for dtype in [
+        dtypes.float32,
+        dtypes.float64,
+        dtypes.bool,
+        dtypes.uint8,
+        dtypes.int8,
+        dtypes.int16,
+        dtypes.int32,
+        dtypes.int64,
+        dtypes.float8_e5m2,
+        dtypes.float8_e4m3fn,
+        dtypes.float8_e4m3fnuz,
+        dtypes.float8_e4m3b11fnuz,
+        dtypes.float8_e5m2fnuz,
+        dtypes.float4_e2m1fn,
+        dtypes.int4,
+        dtypes.uint4,
+        dtypes.int2,
+        dtypes.uint2,
+    ]:
+      self.assertIs(dtype.real_dtype, dtype)
+    self.assertIs(dtypes.complex64.real_dtype, dtypes.float32)
+    self.assertIs(dtypes.complex128.real_dtype, dtypes.float64)
+
+  def testStringConversion(self):
+    self.assertIs(dtypes.float32, dtypes.as_dtype("float32"))
+    self.assertIs(dtypes.float64, dtypes.as_dtype("float64"))
+    self.assertIs(dtypes.int32, dtypes.as_dtype("int32"))
+    self.assertIs(dtypes.uint8, dtypes.as_dtype("uint8"))
+    self.assertIs(dtypes.uint16, dtypes.as_dtype("uint16"))
+    self.assertIs(dtypes.int16, dtypes.as_dtype("int16"))
+    self.assertIs(dtypes.int8, dtypes.as_dtype("int8"))
+    self.assertIs(dtypes.string, dtypes.as_dtype("string"))
+    self.assertIs(dtypes.complex64, dtypes.as_dtype("complex64"))
+    self.assertIs(dtypes.complex128, dtypes.as_dtype("complex128"))
+    self.assertIs(dtypes.int64, dtypes.as_dtype("int64"))
+    self.assertIs(dtypes.bool, dtypes.as_dtype("bool"))
+    self.assertIs(dtypes.qint8, dtypes.as_dtype("qint8"))
+    self.assertIs(dtypes.quint8, dtypes.as_dtype("quint8"))
+    self.assertIs(dtypes.qint32, dtypes.as_dtype("qint32"))
+    self.assertIs(dtypes.bfloat16, dtypes.as_dtype("bfloat16"))
+    self.assertIs(dtypes.float8_e5m2, dtypes.as_dtype("float8_e5m2"))
+    self.assertIs(dtypes.float8_e4m3fn, dtypes.as_dtype("float8_e4m3fn"))
+    self.assertIs(dtypes.float8_e4m3fnuz, dtypes.as_dtype("float8_e4m3fnuz"))
+    self.assertIs(
+        dtypes.float8_e4m3b11fnuz, dtypes.as_dtype("float8_e4m3b11fnuz")
+    )
+    self.assertIs(dtypes.float8_e5m2fnuz, dtypes.as_dtype("float8_e5m2fnuz"))
+    self.assertIs(dtypes.float4_e2m1fn, dtypes.as_dtype("float4_e2m1fn"))
+    self.assertIs(dtypes.int4, dtypes.as_dtype("int4"))
+    self.assertIs(dtypes.uint4, dtypes.as_dtype("uint4"))
+    self.assertIs(dtypes.int2, dtypes.as_dtype("int2"))
+    self.assertIs(dtypes.uint2, dtypes.as_dtype("uint2"))
+    self.assertIs(dtypes.float32_ref, dtypes.as_dtype("float32_ref"))
+    self.assertIs(dtypes.float64_ref, dtypes.as_dtype("float64_ref"))
+    self.assertIs(dtypes.int32_ref, dtypes.as_dtype("int32_ref"))
+    self.assertIs(dtypes.uint8_ref, dtypes.as_dtype("uint8_ref"))
+    self.assertIs(dtypes.int16_ref, dtypes.as_dtype("int16_ref"))
+    self.assertIs(dtypes.int8_ref, dtypes.as_dtype("int8_ref"))
+    self.assertIs(dtypes.string_ref, dtypes.as_dtype("string_ref"))
+    self.assertIs(dtypes.complex64_ref, dtypes.as_dtype("complex64_ref"))
+    self.assertIs(dtypes.complex128_ref, dtypes.as_dtype("complex128_ref"))
+    self.assertIs(dtypes.int64_ref, dtypes.as_dtype("int64_ref"))
+    self.assertIs(dtypes.bool_ref, dtypes.as_dtype("bool_ref"))
+    self.assertIs(dtypes.qint8_ref, dtypes.as_dtype("qint8_ref"))
+    self.assertIs(dtypes.quint8_ref, dtypes.as_dtype("quint8_ref"))
+    self.assertIs(dtypes.qint32_ref, dtypes.as_dtype("qint32_ref"))
+    self.assertIs(dtypes.bfloat16_ref, dtypes.as_dtype("bfloat16_ref"))
+    self.assertIs(dtypes.float8_e5m2_ref, dtypes.as_dtype("float8_e5m2_ref"))
+    self.assertIs(
+        dtypes.float8_e4m3fn_ref, dtypes.as_dtype("float8_e4m3fn_ref")
+    )
+    self.assertIs(
+        dtypes.float4_e2m1fn_ref, dtypes.as_dtype("float4_e2m1fn_ref")
+    )
+    self.assertIs(dtypes.int4_ref, dtypes.as_dtype("int4_ref"))
+    self.assertIs(dtypes.uint4_ref, dtypes.as_dtype("uint4_ref"))
+    self.assertIs(dtypes.int2_ref, dtypes.as_dtype("int2_ref"))
+    self.assertIs(dtypes.uint2_ref, dtypes.as_dtype("uint2_ref"))
+    with self.assertRaises(TypeError):
+      dtypes.as_dtype("not_a_type")
+
+  def testDTypesHaveUniqueNames(self):
+    dtypez = []
+    names = set()
+    for datatype_enum in types_pb2.DataType.values():
+      if datatype_enum == types_pb2.DT_INVALID:
+        continue
+      dtype = dtypes.as_dtype(datatype_enum)
+      dtypez.append(dtype)
+      names.add(dtype.name)
+    self.assertEqual(len(dtypez), len(names))
+
+  def testIsInteger(self):
+    self.assertEqual(dtypes.as_dtype("int8").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("int16").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("int32").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("int64").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("uint8").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("uint16").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("complex64").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("complex128").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("double").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("string").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("bool").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("bfloat16").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fn").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fnuz").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3b11fnuz").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2fnuz").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("float4_e2m1fn").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("int4").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("uint4").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("int2").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("uint2").is_integer, True)
+    self.assertEqual(dtypes.as_dtype("qint8").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("qint16").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("qint32").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("quint8").is_integer, False)
+    self.assertEqual(dtypes.as_dtype("quint16").is_integer, False)
+
+  def testIsFloating(self):
+    self.assertEqual(dtypes.as_dtype("int8").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("int16").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("int32").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("int64").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("uint8").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("uint16").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("complex64").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("complex128").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("float32").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float64").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("string").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("bool").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("bfloat16").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fn").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fnuz").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3b11fnuz").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2fnuz").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("float4_e2m1fn").is_floating, True)
+    self.assertEqual(dtypes.as_dtype("int4").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("uint4").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("int2").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("uint2").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("qint8").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("qint16").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("qint32").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("quint8").is_floating, False)
+    self.assertEqual(dtypes.as_dtype("quint16").is_floating, False)
+
+  def testIsComplex(self):
+    self.assertEqual(dtypes.as_dtype("int8").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("int16").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("int32").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("int64").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("uint8").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("uint16").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("complex64").is_complex, True)
+    self.assertEqual(dtypes.as_dtype("complex128").is_complex, True)
+    self.assertEqual(dtypes.as_dtype("float32").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float64").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("string").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("bool").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("bfloat16").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fn").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fnuz").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3b11fnuz").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2fnuz").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("float4_e2m1fn").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("int4").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("uint4").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("int2").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("uint2").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("qint8").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("qint16").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("qint32").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("quint8").is_complex, False)
+    self.assertEqual(dtypes.as_dtype("quint16").is_complex, False)
+
+  def testIsUnsigned(self):
+    self.assertEqual(dtypes.as_dtype("int8").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("int16").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("int32").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("int64").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("uint8").is_unsigned, True)
+    self.assertEqual(dtypes.as_dtype("uint16").is_unsigned, True)
+    self.assertEqual(dtypes.as_dtype("float32").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float64").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("bool").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("string").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("complex64").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("complex128").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("bfloat16").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fn").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3fnuz").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float8_e4m3b11fnuz").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float8_e5m2fnuz").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("float4_e2m1fn").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("int4").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("uint4").is_unsigned, True)
+    self.assertEqual(dtypes.as_dtype("int2").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("uint2").is_unsigned, True)
+    self.assertEqual(dtypes.as_dtype("qint8").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("qint16").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("qint32").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("quint8").is_unsigned, False)
+    self.assertEqual(dtypes.as_dtype("quint16").is_unsigned, False)
+
+  def testMinMax(self):
+    # make sure min/max evaluates for all data types that have min/max
+    for datatype_enum in types_pb2.DataType.values():
+      if not _is_numeric_dtype_enum(datatype_enum):
+        continue
+      dtype = dtypes.as_dtype(datatype_enum)
+      numpy_dtype = dtype.as_numpy_dtype
+
+      # ignore types for which there are no minimum/maximum (or we cannot
+      # compute it, such as for the q* types)
+      if (dtype.is_quantized or dtype.base_dtype == dtypes.bool or
+          dtype.base_dtype == dtypes.string or
+          dtype.base_dtype == dtypes.complex64 or
+          dtype.base_dtype == dtypes.complex128):
+        continue
+
+      print("%s: %s - %s" % (dtype, dtype.min, dtype.max))
+
+      # check some values that are known
+      if numpy_dtype == np.bool_:
+        self.assertEqual(dtype.min, 0)
+        self.assertEqual(dtype.max, 1)
+      if numpy_dtype == np.int8:
+        self.assertEqual(dtype.min, -128)
+        self.assertEqual(dtype.max, 127)
+      if numpy_dtype == np.int16:
+        self.assertEqual(dtype.min, -32768)
+        self.assertEqual(dtype.max, 32767)
+      if numpy_dtype == np.int32:
+        self.assertEqual(dtype.min, -2147483648)
+        self.assertEqual(dtype.max, 2147483647)
+      if numpy_dtype == np.int64:
+        self.assertEqual(dtype.min, -9223372036854775808)
+        self.assertEqual(dtype.max, 9223372036854775807)
+      if numpy_dtype == np.uint8:
+        self.assertEqual(dtype.min, 0)
+        self.assertEqual(dtype.max, 255)
+      if numpy_dtype == np.uint16:
+        if dtype == dtypes.uint16:
+          self.assertEqual(dtype.min, 0)
+          self.assertEqual(dtype.max, 65535)
+        elif dtype == dtypes.bfloat16:
+          self.assertEqual(dtype.min, 0)
+          self.assertEqual(dtype.max, 4294967295)
+      if numpy_dtype == np.uint32:
+        self.assertEqual(dtype.min, 0)
+        self.assertEqual(dtype.max, 4294967295)
+      if numpy_dtype == np.uint64:
+        self.assertEqual(dtype.min, 0)
+        self.assertEqual(dtype.max, 18446744073709551615)
+      if numpy_dtype in (np.float16, np.float32, np.float64):
+        self.assertEqual(dtype.min, np.finfo(numpy_dtype).min)
+        self.assertEqual(dtype.max, np.finfo(numpy_dtype).max)
+      if numpy_dtype == dtypes.bfloat16.as_numpy_dtype:
+        self.assertEqual(dtype.min, float.fromhex("-0x1.FEp127"))
+        self.assertEqual(dtype.max, float.fromhex("0x1.FEp127"))
+      if numpy_dtype == dtypes.float8_e5m2.as_numpy_dtype:
+        self.assertEqual(dtype.min, -57344.0)
+        self.assertEqual(dtype.max, 57344.0)
+      if numpy_dtype == dtypes.float8_e4m3fn.as_numpy_dtype:
+        self.assertEqual(dtype.min, -448.0)
+        self.assertEqual(dtype.max, 448.0)
+      if numpy_dtype == dtypes.float8_e4m3fnuz.as_numpy_dtype:
+        self.assertEqual(dtype.min, -240.0)
+        self.assertEqual(dtype.max, 240.0)
+      if numpy_dtype == dtypes.float8_e4m3b11fnuz.as_numpy_dtype:
+        self.assertEqual(dtype.min, -30.0)
+        self.assertEqual(dtype.max, 30.0)
+      if numpy_dtype == dtypes.float8_e5m2fnuz.as_numpy_dtype:
+        self.assertEqual(dtype.min, -57344.0)
+        self.assertEqual(dtype.max, 57344.0)
+      if numpy_dtype == dtypes.float4_e2m1fn.as_numpy_dtype:
+        self.assertEqual(dtype.min, -6)
+        self.assertEqual(dtype.max, 6)
+      if numpy_dtype == dtypes.int4.as_numpy_dtype:
+        self.assertEqual(dtype.min, -8)
+        self.assertEqual(dtype.max, 7)
+      if numpy_dtype == dtypes.uint4.as_numpy_dtype:
+        self.assertEqual(dtype.min, 0)
+        self.assertEqual(dtype.max, 15)
+      if numpy_dtype == dtypes.int2.as_numpy_dtype:
+        self.assertEqual(dtype.min, -2)
+        self.assertEqual(dtype.max, 1)
+      if numpy_dtype == dtypes.uint2.as_numpy_dtype:
+        self.assertEqual(dtype.min, 0)
+        self.assertEqual(dtype.max, 3)
+
+  def testLimitsUndefinedError(self):
+    with self.assertRaises(ValueError):
+      dtypes.string.limits()
+
+  def testRepr(self):
+    self.skipTest("b/142725777")
+    for enum, name in dtypes._TYPE_TO_STRING.items():
+      if enum > 100:
+        continue
+      dtype = dtypes.DType(enum)
+      self.assertEqual(repr(dtype), "tf." + name)
+      import tensorflow as tf
+      dtype2 = eval(repr(dtype))
+      self.assertEqual(type(dtype2), dtypes.DType)
+      self.assertEqual(dtype, dtype2)
+
+  def testEqWithNonTFTypes(self):
+    self.assertNotEqual(dtypes.int32, int)
+    self.assertNotEqual(dtypes.float64, 2.1)
+
+  def testPythonLongConversion(self):
+    self.assertIs(dtypes.int64, dtypes.as_dtype(np.array(2**32).dtype))
+
+  def testPythonTypesConversion(self):
+    self.assertIs(dtypes.float32, dtypes.as_dtype(float))
+    self.assertIs(dtypes.bool, dtypes.as_dtype(bool))
+
+  def testReduce(self):
+    for enum in dtypes._TYPE_TO_STRING:
+      dtype = dtypes.DType(enum)
+      ctor, args = dtype.__reduce__()
+      self.assertEqual(ctor, dtypes.as_dtype)
+      self.assertEqual(args, (dtype.name,))
+      reconstructed = ctor(*args)
+      self.assertEqual(reconstructed, dtype)
+
+  def testAsDtypeInvalidArgument(self):
+    with self.assertRaises(TypeError):
+      dtypes.as_dtype((dtypes.int32, dtypes.float32))
+
+  def testAsDtypeReturnsInternedVersion(self):
+    dt = dtypes.DType(types_pb2.DT_VARIANT)
+    self.assertIs(dtypes.as_dtype(dt), dtypes.variant)
+
+  def testDTypeSubtypes(self):
+    self.assertTrue(dtypes.string.is_subtype_of(dtypes.string))
+    self.assertFalse(dtypes.string.is_subtype_of(dtypes.uint32))
+    self.assertTrue(dtypes.uint64.is_subtype_of(dtypes.uint64))
+
+  def testDTypeSupertypes(self):
+    self.assertEqual(dtypes.string,
+                     dtypes.string.most_specific_common_supertype([]))
+    self.assertEqual(
+        dtypes.string,
+        dtypes.string.most_specific_common_supertype([dtypes.string]))
+    self.assertIsNone(
+        dtypes.string.most_specific_common_supertype([dtypes.uint32]))
+
+  @parameterized.parameters(*tuple(dtype for dtype in dtypes.TF_VALUE_DTYPES))
+  def testDTypeSerialization(self, dtype):
+    self.assertEqual(trace_type.deserialize(trace_type.serialize(dtype)), dtype)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/framework/error_interpolation.py

@@ -0,0 +1,438 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Function for interpolating formatted errors from the TensorFlow runtime.
+
+Exposes the function `interpolate` to interpolate messages with tags of the form
+{{type name}}.
+"""
+
+import collections
+import os
+import re
+import site
+import traceback
+
+from tensorflow.python.util import tf_stack
+
+_NAME_REGEX = r"[A-Za-z0-9_.][A-Za-z0-9_.\-/]*?"
+_TAG_REGEX = fr"{{{{(?P<type>{_NAME_REGEX}) (?P<name>{_NAME_REGEX})}}}}"
+_INTERPOLATION_REGEX = fr"(?P<sep>.*?)(?P<tag>{_TAG_REGEX})"
+_INTERPOLATION_PATTERN = re.compile(_INTERPOLATION_REGEX, re.DOTALL)
+
+_ParseTag = collections.namedtuple("_ParseTag", ["type", "name"])
+
+
+# Remove the last three path components from this module's file (i.e.
+# python/framework/error_interpolation.py) so that we have an absolute path
+# prefix to the root of the installation.
+_FRAMEWORK_COMMON_PREFIX = os.path.dirname(
+    os.path.dirname(os.path.dirname(__file__)))
+
+# Sub-directories under the common prefix that are considered part of the
+# framework.
+# Note that keras code lives outside of tensorflow directory, we need to walk
+# up the directory tree and find it.
+_FRAMEWORK_PATH_PREFIXES = [
+    os.path.join(_FRAMEWORK_COMMON_PREFIX, "python") + os.sep,
+    os.path.join(_FRAMEWORK_COMMON_PREFIX, "contrib") + os.sep,
+    os.path.join(os.path.dirname(_FRAMEWORK_COMMON_PREFIX),
+                 "py", "keras") + os.sep,
+]
+
+# Patterns of filename patterns that should be considered internal to
+# the TensorFlow framework.
+_FRAMEWORK_FILENAME_PATTERNS = [
+    re.compile(r"<embedded"),
+]
+
+# This is for OSS keras, since the package is load from local python env,
+# but we don't know exactly where it is installed. Matching to keyword
+# "keras".
+try:
+  _FRAMEWORK_PATH_PREFIXES.extend([
+      os.path.join(package_path, "keras") + os.sep
+      for package_path in site.getsitepackages() + [site.getusersitepackages()]
+  ])
+except AttributeError:
+  # if site.getsitepackages is not available somehow, we just use the "keras" as
+  # the keyword to do the match.
+  _FRAMEWORK_FILENAME_PATTERNS.append(re.compile(r"keras"))
+
+# Patterns of filename patterns that should be considered external to
+# TensorFlow regardless of framework prefix match.
+_EXTERNAL_FILENAME_PATTERNS = [
+    # Explicitly treat test frames as not part of the framework.
+    re.compile(r"_test\.py$"),
+]
+
+
+def parse_message(message):
+  """Extract function tags and node tags from a message.
+
+  Tags are named tuples representing the string {{type name}}. For example,
+  in "123{{node Foo}}456{{function_node Bar}}789", there are two tags: a node
+  tag and a function tag.
+
+  Args:
+    message: An error message, possibly from an OpError.
+
+  Returns:
+    A tuple containing the original message with function nodes stripped,
+    function tags, and node tags.
+
+    For example, if message is "123{{node Foo}}456{{function_node Bar}}789"
+    then this function returns ("123{{node Foo}}456789",
+    [_ParseTag("function_node", "Bar")], [_ParseTag("node", "Foo")]).
+  """
+  error_message = []
+  func_tags = []
+  node_tags = []
+  pos = 0
+  for match in re.finditer(_INTERPOLATION_PATTERN, message):
+    parsed_tag = _ParseTag(match.group("type"), match.group("name"))
+    if parsed_tag.type == "function_node":
+      error_message.append(match.group("sep"))
+      func_tags.append(parsed_tag)
+    else:
+      error_message.append(match.group())
+      node_tags.append(parsed_tag)
+    pos = match.end()
+  error_message.append(message[pos:])
+  return "".join(error_message), func_tags, node_tags
+
+
+def _compute_device_summary_from_list(name, device_assignment_list, prefix=""):
+  """Return a summary of an op's device function stack.
+
+  Args:
+    name: The name of the op.
+    device_assignment_list: The op._device_assignments list.
+    prefix:  An optional string prefix used before each line of the multi-
+        line string returned by this function.
+
+  Returns:
+    A multi-line string similar to:
+        Device assignments active during op 'foo' creation:
+          with tf.device(/cpu:0): <test_1.py:27>
+          with tf.device(some_func<foo.py, 123>): <test_2.py:38>
+    The first line will have no padding to its left by default.  Subsequent
+    lines will have two spaces of left-padding.  Use the prefix argument
+    to increase indentation.
+  """
+  if not device_assignment_list:
+    message = "No device assignments were active during op '%s' creation."
+    message %= name
+    return prefix + message
+
+  str_list = []
+  str_list.append(
+      "%sDevice assignments active during op '%s' creation:" % (prefix, name))
+
+  for traceable_obj in device_assignment_list:
+    location_summary = "<{file}:{line}>".format(
+        file=traceable_obj.filename, line=traceable_obj.lineno)
+    subs = {
+        "prefix": prefix,
+        "indent": "  ",
+        "dev_name": traceable_obj.obj,
+        "loc": location_summary,
+    }
+    str_list.append(
+        "{prefix}{indent}with tf.device({dev_name}): {loc}".format(**subs))
+
+  return "\n".join(str_list)
+
+
+def _compute_device_assignment_summary_from_op(op, prefix=""):
+  # pylint: disable=protected-access
+  return _compute_device_summary_from_list(op.name, op._device_assignments,
+                                           prefix)
+  # pylint: enable=protected-access
+
+
+def _compute_colocation_summary_from_dict(name, colocation_dict, prefix=""):
+  """Return a summary of an op's colocation stack.
+
+  Args:
+    name: The op name.
+    colocation_dict: The op._colocation_dict.
+    prefix:  An optional string prefix used before each line of the multi-
+        line string returned by this function.
+
+  Returns:
+    A multi-line string similar to:
+        Node-device colocations active during op creation:
+          with tf.compat.v1.colocate_with(test_node_1): <test_1.py:27>
+          with tf.compat.v1.colocate_with(test_node_2): <test_2.py:38>
+    The first line will have no padding to its left by default.  Subsequent
+    lines will have two spaces of left-padding.  Use the prefix argument
+    to increase indentation.
+  """
+  if not colocation_dict:
+    message = "No node-device colocations were active during op '%s' creation."
+    message %= name
+    return prefix + message
+
+  str_list = []
+  str_list.append("%sNode-device colocations active during op '%s' creation:" %
+                  (prefix, name))
+
+  for coloc_name, location in colocation_dict.items():
+    location_summary = "<{file}:{line}>".format(
+        file=location.filename, line=location.lineno)
+    subs = {
+        "prefix": prefix,
+        "indent": "  ",
+        "name": coloc_name,
+        "loc": location_summary,
+    }
+    str_list.append(
+        "{prefix}{indent}with tf.colocate_with({name}): {loc}".format(**subs))
+
+  return "\n".join(str_list)
+
+
+def _compute_colocation_summary_from_op(op, prefix=""):
+  """Fetch colocation file, line, and nesting and return a summary string."""
+  # pylint: disable=protected-access
+  return _compute_colocation_summary_from_dict(op.name, op._colocation_dict,
+                                               prefix)
+  # pylint: enable=protected-access
+
+
+def _is_framework_filename(filename):
+  """Returns whether a filename should be considered a part of the framework.
+
+  A file is part of the framework if it does not match a pattern in
+  _EXTERNAL_FILENAME_PATTERNS and it either matches a pattern in
+  _FRAMEWORK_FILENAME_PATTERNS or starts with a _FRAMEWORK_PATH_PREFIXES prefix.
+
+  Args:
+    filename: A filename string.
+
+  Returns:
+    Whether the filename should be considered to be internal to the
+    TensorFlow framework for the purposes of reporting errors.
+  """
+  for pattern in _EXTERNAL_FILENAME_PATTERNS:
+    if pattern.search(filename):
+      return False
+  for pattern in _FRAMEWORK_FILENAME_PATTERNS:
+    if pattern.search(filename):
+      return True
+  for prefix in _FRAMEWORK_PATH_PREFIXES:
+    if filename.startswith(prefix):
+      return True
+  return False
+
+
+def _find_index_of_defining_frame(tb):
+  """Return index in op.traceback with first 'useful' frame.
+
+  This method reads through the stack stored in op.traceback looking for the
+  innermost frame which (hopefully) belongs to the caller.  It accomplishes this
+  by rejecting frames deemed to be part of the TensorFlow framework (by
+  pattern matching the filename).
+
+  Args:
+    tb: A list of traceback frames (as from Operation.traceback).
+
+  Returns:
+    Integer index into op.traceback where the first non-TF file was found
+    (innermost to outermost), or 0 (for the outermost stack frame) if all files
+    came from TensorFlow.
+  """
+  # Index 0 of traceback is the outermost frame.
+  size = len(tb)
+  filenames = [frame.filename for frame in tb]
+  # We process the filenames from the innermost frame to outermost.
+  for idx, filename in enumerate(reversed(filenames)):
+    is_framework = _is_framework_filename(filename)
+    if not is_framework:
+      # Consider this to be the defining frame.
+      return size - idx - 1
+  return 0
+
+
+# TODO(feyu): follow up with users of this function (saved model)
+# to see what 'useful' means and whether we can obliviate this.
+def _compute_useful_frames(tb, num):
+  """Return a list of frames, which form a 'useful' stack.
+
+  Starting from the defining frame to the outermost one, this method computes
+  the contiguous portion of the 'useful' stack trace and returns the selected
+  frames.
+
+  Args:
+    tb: A list of traceback frames (as from Operation.traceback).
+    num: total number of frames to return.
+
+  Returns:
+    A list of frames.
+  """
+  defining_frame_index = _find_index_of_defining_frame(tb)
+  # The stack trace is collected from two lines before the defining frame in the
+  # model file to the outermost with `num` frames at most. These two extra lines
+  # are included from the TensorFlow library to give the context which node is
+  # defined.
+  innermost_excluded = min(defining_frame_index + 2 + 1, len(tb))
+  outermost_included = max(innermost_excluded - num, 0)
+  return tb[outermost_included:innermost_excluded]
+
+
+def create_graph_debug_info_def(func_named_operations):
+  """Construct and returns a `GraphDebugInfo` protocol buffer.
+
+  Args:
+    func_named_operations: An iterable of (func_name, op.Operation) tuples
+      where the Operation instances have a _traceback members. The func_name
+      should be the empty string for operations in the top-level Graph.
+
+  Returns:
+    GraphDebugInfo protocol buffer.
+
+  Raises:
+    TypeError: If the arguments are not of the correct proto buffer type.
+  """
+  builder = tf_stack.GraphDebugInfoBuilder()
+  for func_name, op in func_named_operations:
+    if op.traceback is None:
+      continue
+    builder.AccumulateStackTrace(
+        func_name, op.name, _compute_useful_frames(op.traceback, 10)
+    )
+
+  return builder.Build()
+
+
+def _compute_field_dict(op):
+  r"""Return a dictionary mapping interpolation tokens to values.
+
+  Args:
+    op: op.Operation object.
+
+  Returns:
+    A dictionary mapping string tokens to string values.  The keys are shown
+    below along with example values.
+    {
+      "file": "tool_utils.py",
+      "lineno": "124",
+      "line": "  source code line",
+      "defined_at": " (defined at tool_utils.py:124)",
+      "colocations":
+          '''Node-device colocations active during op creation:
+               with tf.compat.v1.colocate_with(test_node_1): <test_1.py:27>
+               with tf.compat.v1.colocate_with(test_node_2): <test_2.py:38>'''
+      "devices":
+          '''Device assignments active during op 'foo' creation:
+               with tf.device(/cpu:0): <test_1.py:27>
+               with tf.device(some_func<foo.py, 123>): <test_2.py:38>'''
+      "devs_and_colocs": A concatenation of colocations and devices, e.g.
+          '''Node-device colocations active during op creation:
+               with tf.compat.v1.colocate_with(test_node_1): <test_1.py:27>
+               with tf.compat.v1.colocate_with(test_node_2): <test_2.py:38>'''
+             Device assignments active during op 'foo' creation:
+               with tf.device(/cpu:0): <test_1.py:27>
+               with tf.device(some_func<foo.py, 123>): <test_2.py:38>'''
+    }
+  """
+  # TODO(xjun): colocation and device info are not displayed. Consider
+  # removing them or using vlog.
+  colocation_summary = _compute_colocation_summary_from_op(op)
+  device_summary = _compute_device_assignment_summary_from_op(op)
+  combined_summary = "\n".join([colocation_summary, device_summary])
+
+  if op.traceback is None:
+    # Some ops synthesized on as part of function or control flow definition
+    # do not have tracebacks.
+    filename = "<unknown>"
+    definition_traceback = ""
+    lineno = 0
+    line = ""
+    defined_at = "<unknown>"
+  else:
+    frame = op.traceback.last_user_frame()
+    filename = frame.filename
+    definition_traceback = traceback.format_list(op.traceback.get_user_frames())
+    lineno = frame.lineno
+    line = frame.line
+    defined_at = f"{filename}:{lineno:d}"
+
+  field_dict = {
+      "colocations": colocation_summary,
+      "devices": device_summary,
+      "devs_and_colocs": combined_summary,
+      "defined_at": defined_at,
+      "file": filename,
+      "lineno": lineno,
+      "line": line,
+      "definition_traceback": definition_traceback,
+  }
+  return field_dict
+
+
+def _build_node_error_message(op):
+  """Returns the formatted error message for the given op.
+
+  Args:
+    op: The node.
+
+  Returns:
+    The formatted error message for the given op with traceback.
+  """
+  node_error_message = [
+      f"Detected at node {op.name!r} defined at (most recent call last):"
+  ]
+  field_dict = _compute_field_dict(op)
+
+  # Add node traceback.
+  for frame in field_dict["definition_traceback"]:
+    if "<embedded" not in frame:
+      node_error_message.extend(
+          [f"  {line}" for line in frame.split("\n") if line.strip()])
+
+  # Add node name.
+  node_error_message.append(f"Node: {op.name!r}")
+
+  return "\n".join(node_error_message)
+
+
+def interpolate_graph(message, graph):
+  """Interpolates an error message.
+
+  The error message can contain tags of form `{{node_type node_name}}`
+  which will be parsed to identify the tf.Graph and op. If the op contains
+  traceback, the traceback will be attached to the error message.
+
+  Args:
+    message: A string to interpolate.
+    graph: ops.Graph object containing all nodes referenced in the error
+        message.
+
+  Returns:
+    The error message string with node definition traceback.
+  """
+  parsed_messaged, _, node_tags = parse_message(message)
+  error_message = ["Graph execution error:", ""]
+  for tag in node_tags:
+    try:
+      op = graph.get_operation_by_name(tag.name)
+    except KeyError:
+      continue
+    else:
+      error_message.append(_build_node_error_message(op))
+
+  error_message.append(parsed_messaged.strip())
+  return "\n".join(error_message)

tensorflow/python/framework/error_interpolation_test.py

@@ -0,0 +1,338 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.errors."""
+
+import collections
+import os
+import re
+
+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 error_interpolation
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import traceable_stack
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import script_ops
+from tensorflow.python.platform import test
+
+
+# A mock for ``tf_stack.FrameSummary``.
+FrameSummary = collections.namedtuple(
+    "StackFrame", ["filename", "lineno", "name", "line"]
+)
+
+# TODO(feyu): convert tests to tf function from graph when appropriate.
+
+
+def _make_frame_with_filename(tb, idx, filename):
+  """Return a copy of an existing stack frame with a new filename."""
+  frame = tb[idx]
+  return FrameSummary(filename, frame.lineno, frame.name, frame.line)
+
+
+def _modify_op_stack_with_filenames(
+    tb, num_user_frames, user_filename, num_inner_tf_frames
+):
+  """Replace traceback with a new traceback using special filenames."""
+  tf_filename = error_interpolation._FRAMEWORK_PATH_PREFIXES[0] + "%d.py"
+  user_filename = os.path.join("%d", "my_favorite_file.py")
+
+  num_requested_frames = num_user_frames + num_inner_tf_frames
+  num_actual_frames = len(tb)
+  num_outer_frames = num_actual_frames - num_requested_frames
+  assert num_requested_frames <= num_actual_frames, "Too few real frames."
+
+  # The op's traceback has outermost frame at index 0.
+  stack = []
+  for idx in range(0, num_outer_frames):
+    stack.append(tb[idx])
+  for idx in range(len(stack), len(stack) + num_user_frames):
+    stack.append(_make_frame_with_filename(tb, idx, user_filename % idx))
+  for idx in range(len(stack), len(stack) + num_inner_tf_frames):
+    stack.append(_make_frame_with_filename(tb, idx, tf_filename % idx))
+  return stack
+
+
+class ComputeDeviceSummaryFromOpTest(test.TestCase):
+
+  def testCorrectFormatWithActiveDeviceAssignments(self):
+    assignments = []
+    assignments.append(
+        traceable_stack.TraceableObject("/cpu:0", filename="hope.py", lineno=24)
+    )
+    assignments.append(
+        traceable_stack.TraceableObject(
+            "/gpu:2", filename="please.py", lineno=42
+        )
+    )
+
+    summary = error_interpolation._compute_device_summary_from_list(
+        "nodename", assignments, prefix="  "
+    )
+
+    self.assertIn("nodename", summary)
+    self.assertIn("tf.device(/cpu:0)", summary)
+    self.assertIn("<hope.py:24>", summary)
+    self.assertIn("tf.device(/gpu:2)", summary)
+    self.assertIn("<please.py:42>", summary)
+
+  def testCorrectFormatWhenNoColocationsWereActive(self):
+    device_assignment_list = []
+    summary = error_interpolation._compute_device_summary_from_list(
+        "nodename", device_assignment_list, prefix="  "
+    )
+    self.assertIn("nodename", summary)
+    self.assertIn("No device assignments", summary)
+
+
+class ComputeColocationSummaryFromOpTest(test.TestCase):
+
+  def testCorrectFormatWithActiveColocations(self):
+    t_obj_1 = traceable_stack.TraceableObject(
+        None, filename="test_1.py", lineno=27
+    )
+    t_obj_2 = traceable_stack.TraceableObject(
+        None, filename="test_2.py", lineno=38
+    )
+    colocation_dict = {
+        "test_node_1": t_obj_1,
+        "test_node_2": t_obj_2,
+    }
+    summary = error_interpolation._compute_colocation_summary_from_dict(
+        "node_name", colocation_dict, prefix="  "
+    )
+    self.assertIn("node_name", summary)
+    self.assertIn("colocate_with(test_node_1)", summary)
+    self.assertIn("<test_1.py:27>", summary)
+    self.assertIn("colocate_with(test_node_2)", summary)
+    self.assertIn("<test_2.py:38>", summary)
+
+  def testCorrectFormatWhenNoColocationsWereActive(self):
+    colocation_dict = {}
+    summary = error_interpolation._compute_colocation_summary_from_dict(
+        "node_name", colocation_dict, prefix="  "
+    )
+    self.assertIn("node_name", summary)
+    self.assertIn("No node-device colocations", summary)
+
+
+class InterpolateFilenamesAndLineNumbersTest(test.TestCase):
+
+  def testFindIndexOfDefiningFrameForOp(self):
+    with ops.Graph().as_default():
+      local_op = constant_op.constant(42).op
+      user_filename = "hope.py"
+      modified_tb = _modify_op_stack_with_filenames(
+          local_op.traceback,
+          num_user_frames=3,
+          user_filename=user_filename,
+          num_inner_tf_frames=5,
+      )
+      idx = error_interpolation._find_index_of_defining_frame(modified_tb)
+      # Expected frame is 6th from the end because there are 5 inner frames with
+      # TF filenames.
+      expected_frame = len(modified_tb) - 6
+      self.assertEqual(expected_frame, idx)
+
+  def testFindIndexOfDefiningFrameForOpReturnsZeroOnError(self):
+    with ops.Graph().as_default():
+      local_op = constant_op.constant(43).op
+      # Ensure all frames look like TF frames.
+      modified_tb = _modify_op_stack_with_filenames(
+          local_op.traceback[:7],  # Truncate stack to known length.
+          num_user_frames=0,
+          user_filename="user_file.py",
+          num_inner_tf_frames=7,
+      )
+      idx = error_interpolation._find_index_of_defining_frame(modified_tb)
+      self.assertEqual(0, idx)
+
+  def testNothingToDo(self):
+    with ops.Graph().as_default():
+      constant_op.constant(1, name="One")
+      normal_string = "This is just a normal string"
+      interpolated_string = error_interpolation.interpolate_graph(
+          normal_string, ops.get_default_graph()
+      )
+      self.assertIn(normal_string, interpolated_string)
+
+  def testOneTagWithAFakeNameResultsInPlaceholders(self):
+    with ops.Graph().as_default():
+      one_tag_string = "{{node MinusOne}}"
+      interpolated_string = error_interpolation.interpolate_graph(
+          one_tag_string, ops.get_default_graph()
+      )
+      self.assertIn(one_tag_string, interpolated_string)
+
+  def testOneTagWithAFakeFunctionTag(self):
+    defined_at = r"defined at.*error_interpolation_test\.py"
+    with ops.Graph().as_default():
+      constant_op.constant(1, name="One")
+      constant_op.constant(2, name="Two")
+      one_tag_with_a_fake_function_tag = "{{function_node fake}}{{node One}}"
+      interpolated_string = error_interpolation.interpolate_graph(
+          one_tag_with_a_fake_function_tag, ops.get_default_graph()
+      )
+      # Fragments the expression to avoid matching the pattern itself.
+      expected_regex = re.compile(rf"node 'One'.*{defined_at}", re.DOTALL)
+      self.assertRegex(interpolated_string, expected_regex)
+      self.assertNotIn("function_node", interpolated_string)
+      self.assertNotIn("node 'Two'", interpolated_string)
+
+  def testTwoTagsNoSeps(self):
+    defined_at = r"defined at.*error_interpolation_test\.py"
+    with ops.Graph().as_default():
+      constant_op.constant(1, name="One")
+      constant_op.constant(2, name="Two")
+      constant_op.constant(3, name="Three")
+      two_tags_no_seps = "{{node One}}{{node Three}}"
+      interpolated_string = error_interpolation.interpolate_graph(
+          two_tags_no_seps, ops.get_default_graph()
+      )
+      # Fragments the expression to avoid matching the pattern itself.
+      expected_regex = re.compile(
+          rf"node 'One'.*{defined_at}.*node 'Three'.*{defined_at}", re.DOTALL
+      )
+      self.assertRegex(interpolated_string, expected_regex)
+
+  def testTwoTagsWithSeps(self):
+    defined_at = r"defined at.*error_interpolation_test\.py"
+    with ops.Graph().as_default():
+      constant_op.constant(1, name="One")
+      constant_op.constant(2, name="Two")
+      constant_op.constant(3, name="Three")
+      two_tags_with_seps = ";;;{{node Two}},,,{{node Three}};;;"
+      interpolated_string = error_interpolation.interpolate_graph(
+          two_tags_with_seps, ops.get_default_graph()
+      )
+      # Fragments the expression to avoid matching the pattern itself.
+      expected_regex = re.compile(
+          rf"node 'Two'.*{defined_at}.*node 'Three'.*{defined_at}", re.DOTALL
+      )
+      self.assertRegex(interpolated_string, expected_regex)
+
+  def testNewLine(self):
+    defined_at = r"defined at.*error_interpolation_test\.py"
+    with ops.Graph().as_default():
+      constant_op.constant(1, name="One")
+      constant_op.constant(2, name="Two")
+      newline = "\n\n;;;{{node One}};;;"
+      interpolated_string = error_interpolation.interpolate_graph(
+          newline, ops.get_default_graph()
+      )
+      expected_regex = re.compile(rf"node 'One'.*{defined_at}", re.DOTALL)
+      self.assertRegex(interpolated_string, expected_regex)
+
+
+class OperationDefinedAtTraceTest(test.TestCase):
+
+  @test_util.run_v2_only
+  def testSimpleCall(self):
+    @def_function.function
+    def func():
+      x = constant_op.constant([[1, 2, 3]])
+      y = script_ops.eager_py_func(lambda: [[1, 2, 3]], (), dtypes.int32)
+      return math_ops.matmul(x, y)
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        re.compile(
+            r"defined at.*" r"in testSimpleCall.*" r"in func", re.DOTALL
+        ),
+    ):
+      func()
+
+  @test_util.run_v2_only
+  def testNestedCall(self):
+    def inner():
+      x = constant_op.constant([[1, 2, 3]])
+      y = script_ops.eager_py_func(lambda: [[1, 2, 3]], (), dtypes.int32)
+      return math_ops.matmul(x, y)
+
+    @def_function.function
+    def func():
+      return inner()
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        re.compile(
+            r"defined at.*" r"in testNestedCall.*" r"in func.*" r"in inner",
+            re.DOTALL,
+        ),
+    ):
+      func()
+
+  @test_util.run_v2_only
+  def testAssert(self):
+    @def_function.function
+    def func():
+      control_flow_assert.Assert(False, [False])
+      return
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        re.compile(r"defined at.*" r"in testAssert.*" r"in func", re.DOTALL),
+    ):
+      func()
+
+  @test_util.run_v2_only
+  def testControlFlow(self):
+    @def_function.function
+    def func():
+      if constant_op.constant(False):
+        return constant_op.constant(1)
+
+      else:
+        x = constant_op.constant([[1, 2, 3]])
+        y = script_ops.eager_py_func(lambda: [[1, 2, 3]], (), dtypes.int32)
+        return math_ops.matmul(x, y)
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        re.compile(
+            r"defined at.*" r"in testControlFlow.*" r"in func", re.DOTALL
+        ),
+    ):
+      func()
+
+
+class IsFrameworkFilenameTest(test.TestCase):
+
+  def testAllowsUnitTests(self):
+    self.assertFalse(
+        error_interpolation._is_framework_filename(
+            error_interpolation._FRAMEWORK_PATH_PREFIXES[0] + "foobar_test.py"
+        )
+    )
+
+  def testFrameworkPythonFile(self):
+    self.assertTrue(
+        error_interpolation._is_framework_filename(error_interpolation.__file__)
+    )
+
+  def testEmbedded(self):
+    self.assertTrue(
+        error_interpolation._is_framework_filename(
+            "<embedded stdlib>/context_lib.py"
+        )
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/framework/errors.py

@@ -0,0 +1,26 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Exception types for TensorFlow errors.
+
+API docstring: tensorflow.errors
+"""
+
+# pylint: disable=unused-import
+from tensorflow.python.framework import errors_impl as _impl
+# pylint: enable=unused-import
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.framework.errors_impl import *
+# pylint: enable=wildcard-import

tensorflow/python/framework/errors_impl.py

@@ -0,0 +1,571 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Exception types for TensorFlow errors."""
+
+import traceback
+import warnings
+
+from tensorflow.core.lib.core import error_codes_pb2
+from tensorflow.python import _pywrap_py_exception_registry
+from tensorflow.python.client import pywrap_tf_session as c_api
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util.tf_export import tf_export
+
+
+class InaccessibleTensorError(ValueError):
+  pass
+
+
+@tf_export("errors.OperatorNotAllowedInGraphError", v1=[])
+class OperatorNotAllowedInGraphError(TypeError):
+  """Raised when an unsupported operator is present in Graph execution.
+
+  For example, using a `tf.Tensor` as a Python `bool` inside a Graph will
+  raise `OperatorNotAllowedInGraphError`. Iterating over values inside a
+  `tf.Tensor` is also not supported in Graph execution.
+
+  Example:
+  >>> @tf.function
+  ... def iterate_over(t):
+  ...   a,b,c = t
+  ...   return a
+  >>>
+  >>> iterate_over(tf.constant([1, 2, 3]))
+  Traceback (most recent call last):
+  ...
+  OperatorNotAllowedInGraphError: ...
+
+  """
+  pass
+
+
+@tf_export("errors.OpError", v1=["errors.OpError", "OpError"])
+@deprecation.deprecated_endpoints("OpError")
+class OpError(Exception):
+  """The base class for TensorFlow exceptions.
+
+  Usually, TensorFlow will raise a more specific subclass of `OpError` from the
+  `tf.errors` module.
+  """
+
+  def __init__(self, node_def, op, message, error_code, *args):
+    """Creates a new `OpError` indicating that a particular op failed.
+
+    Args:
+      node_def: The `node_def_pb2.NodeDef` proto representing the op that
+        failed, if known; otherwise None.
+      op: The `ops.Operation` that failed, if known; otherwise None. During
+        eager execution, this field is always `None`.
+      message: The message string describing the failure.
+      error_code: The `error_codes_pb2.Code` describing the error.
+      *args: If not empty, it should contain a dictionary describing details
+        about the error. This argument is inspired by Abseil payloads:
+        https://github.com/abseil/abseil-cpp/blob/master/absl/status/status.h
+    """
+    super(OpError, self).__init__()
+    self._node_def = node_def
+    self._op = op
+    self._message = message
+    self._error_code = error_code
+    if args:
+      self._experimental_payloads = args[0]
+    else:
+      self._experimental_payloads = {}
+
+  def __reduce__(self):
+    # Allow the subclasses to accept less arguments in their __init__.
+    init_argspec = tf_inspect.getargspec(self.__class__.__init__)
+    args = tuple(getattr(self, arg) for arg in init_argspec.args[1:])
+    return self.__class__, args
+
+  @property
+  def message(self):
+    """The error message that describes the error."""
+    return self._message
+
+  @property
+  def op(self):
+    """The operation that failed, if known.
+
+    *N.B.* If the failed op was synthesized at runtime, e.g. a `Send`
+    or `Recv` op, there will be no corresponding
+    `tf.Operation`
+    object.  In that case, this will return `None`, and you should
+    instead use the `tf.errors.OpError.node_def` to
+    discover information about the op.
+
+    Returns:
+      The `Operation` that failed, or None.
+    """
+    return self._op
+
+  @property
+  def error_code(self):
+    """The integer error code that describes the error."""
+    return self._error_code
+
+  @property
+  def node_def(self):
+    """The `NodeDef` proto representing the op that failed."""
+    return self._node_def
+
+  @property
+  def experimental_payloads(self):
+    """A dictionary describing the details of the error."""
+    return self._experimental_payloads
+
+  def __str__(self):
+    if self._op is not None:
+      output = [
+          "%s\n\nOriginal stack trace for %r:\n" % (
+              self.message,
+              self._op.name,
+          )
+      ]
+      curr_traceback_list = traceback.format_list(self._op.traceback or [])
+      output.extend(curr_traceback_list)
+      # pylint: disable=protected-access
+      original_op = self._op._original_op
+      # pylint: enable=protected-access
+      while original_op is not None:
+        output.append(
+            "\n...which was originally created as op %r, defined at:\n" %
+            (original_op.name,))
+        prev_traceback_list = curr_traceback_list
+        curr_traceback_list = traceback.format_list(original_op.traceback or [])
+
+        # Attempt to elide large common subsequences of the subsequent
+        # stack traces.
+        #
+        # TODO(mrry): Consider computing the actual longest common subsequence.
+        is_eliding = False
+        elide_count = 0
+        last_elided_line = None
+        for line, line_in_prev in zip(curr_traceback_list, prev_traceback_list):
+          if line == line_in_prev:
+            if is_eliding:
+              elide_count += 1
+              last_elided_line = line
+            else:
+              output.append(line)
+              is_eliding = True
+              elide_count = 0
+          else:
+            if is_eliding:
+              if elide_count > 0:
+                output.extend([
+                    "[elided %d identical lines from previous traceback]\n" %
+                    (elide_count - 1,), last_elided_line
+                ])
+              is_eliding = False
+            output.extend(line)
+
+        # pylint: disable=protected-access
+        original_op = original_op._original_op
+        # pylint: enable=protected-access
+      return "".join(output)
+    else:
+      return self.message
+
+
+OK = error_codes_pb2.OK
+tf_export("errors.OK").export_constant(__name__, "OK")
+CANCELLED = error_codes_pb2.CANCELLED
+tf_export("errors.CANCELLED").export_constant(__name__, "CANCELLED")
+UNKNOWN = error_codes_pb2.UNKNOWN
+tf_export("errors.UNKNOWN").export_constant(__name__, "UNKNOWN")
+INVALID_ARGUMENT = error_codes_pb2.INVALID_ARGUMENT
+tf_export("errors.INVALID_ARGUMENT").export_constant(__name__,
+                                                     "INVALID_ARGUMENT")
+DEADLINE_EXCEEDED = error_codes_pb2.DEADLINE_EXCEEDED
+tf_export("errors.DEADLINE_EXCEEDED").export_constant(__name__,
+                                                      "DEADLINE_EXCEEDED")
+NOT_FOUND = error_codes_pb2.NOT_FOUND
+tf_export("errors.NOT_FOUND").export_constant(__name__, "NOT_FOUND")
+ALREADY_EXISTS = error_codes_pb2.ALREADY_EXISTS
+tf_export("errors.ALREADY_EXISTS").export_constant(__name__, "ALREADY_EXISTS")
+PERMISSION_DENIED = error_codes_pb2.PERMISSION_DENIED
+tf_export("errors.PERMISSION_DENIED").export_constant(__name__,
+                                                      "PERMISSION_DENIED")
+UNAUTHENTICATED = error_codes_pb2.UNAUTHENTICATED
+tf_export("errors.UNAUTHENTICATED").export_constant(__name__, "UNAUTHENTICATED")
+RESOURCE_EXHAUSTED = error_codes_pb2.RESOURCE_EXHAUSTED
+tf_export("errors.RESOURCE_EXHAUSTED").export_constant(__name__,
+                                                       "RESOURCE_EXHAUSTED")
+FAILED_PRECONDITION = error_codes_pb2.FAILED_PRECONDITION
+tf_export("errors.FAILED_PRECONDITION").export_constant(__name__,
+                                                        "FAILED_PRECONDITION")
+ABORTED = error_codes_pb2.ABORTED
+tf_export("errors.ABORTED").export_constant(__name__, "ABORTED")
+OUT_OF_RANGE = error_codes_pb2.OUT_OF_RANGE
+tf_export("errors.OUT_OF_RANGE").export_constant(__name__, "OUT_OF_RANGE")
+UNIMPLEMENTED = error_codes_pb2.UNIMPLEMENTED
+tf_export("errors.UNIMPLEMENTED").export_constant(__name__, "UNIMPLEMENTED")
+INTERNAL = error_codes_pb2.INTERNAL
+tf_export("errors.INTERNAL").export_constant(__name__, "INTERNAL")
+UNAVAILABLE = error_codes_pb2.UNAVAILABLE
+tf_export("errors.UNAVAILABLE").export_constant(__name__, "UNAVAILABLE")
+DATA_LOSS = error_codes_pb2.DATA_LOSS
+tf_export("errors.DATA_LOSS").export_constant(__name__, "DATA_LOSS")
+
+
+@tf_export("errors.CancelledError")
+class CancelledError(OpError):
+  """Raised when an operation is cancelled.
+
+  For example, a long-running operation e.g.`tf.queue.QueueBase.enqueue`, or a
+  `tf.function` call may be cancelled by either running another operation e.g.
+  `tf.queue.QueueBase.close` or a remote worker failure.
+
+  This long-running operation will fail by raising `CancelledError`.
+
+  Example:
+  >>> q = tf.queue.FIFOQueue(10, tf.float32, ((),))
+  >>> q.enqueue((10.0,))
+  >>> q.close()
+  >>> q.enqueue((10.0,))
+  Traceback (most recent call last):
+    ...
+  CancelledError: ...
+
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `CancelledError`."""
+    super(CancelledError, self).__init__(node_def, op, message, CANCELLED,
+                                         *args)
+
+
+@tf_export("errors.UnknownError")
+class UnknownError(OpError):
+  """Unknown error.
+
+  An example of where this error may be returned is if a Status value
+  received from another address space belongs to an error-space that
+  is not known to this address space. Also, errors raised by APIs that
+  do not return enough error information may be converted to this
+  error.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `UnknownError`."""
+    super(UnknownError, self).__init__(node_def, op, message, UNKNOWN, *args)
+
+
+@tf_export("errors.InvalidArgumentError")
+class InvalidArgumentError(OpError):
+  """Raised when an operation receives an invalid argument.
+
+  This error is typically raised when an op receives mismatched arguments.
+
+  Example:
+
+  >>> tf.reshape([1, 2, 3], (2,))
+  Traceback (most recent call last):
+     ...
+  InvalidArgumentError: ...
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `InvalidArgumentError`."""
+    super(InvalidArgumentError, self).__init__(node_def, op, message,
+                                               INVALID_ARGUMENT, *args)
+
+
+@tf_export("errors.DeadlineExceededError")
+class DeadlineExceededError(OpError):
+  """Raised when a deadline expires before an operation could complete.
+
+  This exception is not currently used.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `DeadlineExceededError`."""
+    super(DeadlineExceededError, self).__init__(node_def, op, message,
+                                                DEADLINE_EXCEEDED, *args)
+
+
+@tf_export("errors.NotFoundError")
+class NotFoundError(OpError):
+  """Raised when a requested entity (e.g., a file or directory) was not found.
+
+  For example, running the
+  `tf.WholeFileReader.read`
+  operation could raise `NotFoundError` if it receives the name of a file that
+  does not exist.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `NotFoundError`."""
+    super(NotFoundError, self).__init__(node_def, op, message, NOT_FOUND, *args)
+
+
+@tf_export("errors.AlreadyExistsError")
+class AlreadyExistsError(OpError):
+  """Raised when an entity that we attempted to create already exists.
+
+  An API raises this this error to avoid overwriting an existing resource,
+  value, etc. Calling a creation API multiple times with the same arguments
+  could raise this error if the creation API is not idempotent.
+
+  For example, running an operation that saves a file
+  (e.g. `tf.saved_model.save`)
+  could potentially raise this exception if an explicit filename for an
+  existing file was passed.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `AlreadyExistsError`."""
+    super(AlreadyExistsError, self).__init__(node_def, op, message,
+                                             ALREADY_EXISTS, *args)
+
+
+@tf_export("errors.PermissionDeniedError")
+class PermissionDeniedError(OpError):
+  """Raised when the caller does not have permission to run an operation.
+
+  For example, running the
+  `tf.WholeFileReader.read`
+  operation could raise `PermissionDeniedError` if it receives the name of a
+  file for which the user does not have the read file permission.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `PermissionDeniedError`."""
+    super(PermissionDeniedError, self).__init__(node_def, op, message,
+                                                PERMISSION_DENIED, *args)
+
+
+@tf_export("errors.UnauthenticatedError")
+class UnauthenticatedError(OpError):
+  """Raised when the request does not have valid authentication credentials.
+
+  This exception is not currently used.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `UnauthenticatedError`."""
+    super(UnauthenticatedError, self).__init__(node_def, op, message,
+                                               UNAUTHENTICATED, *args)
+
+
+@tf_export("errors.ResourceExhaustedError")
+class ResourceExhaustedError(OpError):
+  """Raised when some resource has been exhausted while running operation.
+
+  For example, this error might be raised if a per-user quota is
+  exhausted, or perhaps the entire file system is out of space. If running into
+  `ResourceExhaustedError` due to out of memory (OOM), try to use smaller batch
+  size or reduce dimension size of model weights.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `ResourceExhaustedError`."""
+    super(ResourceExhaustedError, self).__init__(node_def, op, message,
+                                                 RESOURCE_EXHAUSTED, *args)
+
+
+@tf_export("errors.FailedPreconditionError")
+class FailedPreconditionError(OpError):
+  """Raised when some prerequisites are not met when running an operation.
+
+  This typically indicates that system is not in state to execute the operation
+  and requires preconditions to be met before successfully executing current
+  operation.
+
+  For example, this exception is commonly raised when running an operation
+  that reads a `tf.Variable` before it has been initialized.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `FailedPreconditionError`."""
+    super(FailedPreconditionError, self).__init__(node_def, op, message,
+                                                  FAILED_PRECONDITION, *args)
+
+
+@tf_export("errors.AbortedError")
+class AbortedError(OpError):
+  """Raised when an operation was aborted, typically due to a concurrent action.
+
+  For example, running a
+  `tf.queue.QueueBase.enqueue`
+  operation may raise `AbortedError` if a
+  `tf.queue.QueueBase.close` operation
+  previously ran.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `AbortedError`."""
+    super(AbortedError, self).__init__(node_def, op, message, ABORTED, *args)
+
+
+@tf_export("errors.OutOfRangeError")
+class OutOfRangeError(OpError):
+  """Raised when an operation iterates past the valid range.
+
+  Unlike `InvalidArgumentError`, this error indicates a problem may be fixed if
+  the system state changes. For example, if a list grows and the operation is
+  now within the valid range. `OutOfRangeError` overlaps with
+  `FailedPreconditionError` and should be preferred as the more specific error
+  when iterating or accessing a range.
+
+  For example, iterating a TF dataset past the last item in the dataset will
+  raise this error.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `OutOfRangeError`."""
+    super(OutOfRangeError, self).__init__(node_def, op, message, OUT_OF_RANGE,
+                                          *args)
+
+
+@tf_export("errors.UnimplementedError")
+class UnimplementedError(OpError):
+  """Raised when an operation has not been implemented.
+
+  Some operations may raise this error when passed otherwise-valid
+  arguments that it does not currently support. For example, running
+  the `tf.nn.max_pool2d` operation
+  would raise this error if pooling was requested on the batch dimension,
+  because this is not yet supported.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `UnimplementedError`."""
+    super(UnimplementedError, self).__init__(node_def, op, message,
+                                             UNIMPLEMENTED, *args)
+
+
+@tf_export("errors.InternalError")
+class InternalError(OpError):
+  """Raised when the system experiences an internal error.
+
+  This exception is raised when some invariant expected by the runtime
+  has been broken. Catching this exception is not recommended.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `InternalError`."""
+    super(InternalError, self).__init__(node_def, op, message, INTERNAL, *args)
+
+
+@tf_export("errors.UnavailableError")
+class UnavailableError(OpError):
+  """Raised when the runtime is currently unavailable.
+
+  This exception is not currently used.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates an `UnavailableError`."""
+    super(UnavailableError, self).__init__(node_def, op, message, UNAVAILABLE,
+                                           *args)
+
+
+@tf_export("errors.DataLossError")
+class DataLossError(OpError):
+  """Raised when unrecoverable data loss or corruption is encountered.
+
+  This could be due to:
+  * A truncated file.
+  * A corrupted file.
+  * Specifying the wrong data format.
+
+  For example, this may be raised by running a
+  `tf.WholeFileReader.read`
+  operation, if the file is truncated while it is being read.
+  """
+
+  def __init__(self, node_def, op, message, *args):
+    """Creates a `DataLossError`."""
+    super(DataLossError, self).__init__(node_def, op, message, DATA_LOSS, *args)
+
+
+_CODE_TO_EXCEPTION_CLASS = {
+    CANCELLED: CancelledError,
+    UNKNOWN: UnknownError,
+    INVALID_ARGUMENT: InvalidArgumentError,
+    DEADLINE_EXCEEDED: DeadlineExceededError,
+    NOT_FOUND: NotFoundError,
+    ALREADY_EXISTS: AlreadyExistsError,
+    PERMISSION_DENIED: PermissionDeniedError,
+    UNAUTHENTICATED: UnauthenticatedError,
+    RESOURCE_EXHAUSTED: ResourceExhaustedError,
+    FAILED_PRECONDITION: FailedPreconditionError,
+    ABORTED: AbortedError,
+    OUT_OF_RANGE: OutOfRangeError,
+    UNIMPLEMENTED: UnimplementedError,
+    INTERNAL: InternalError,
+    UNAVAILABLE: UnavailableError,
+    DATA_LOSS: DataLossError,
+}
+
+_pywrap_py_exception_registry.PyExceptionRegistry_Init(_CODE_TO_EXCEPTION_CLASS)
+
+_EXCEPTION_CLASS_TO_CODE = {
+    class_: code for code, class_ in _CODE_TO_EXCEPTION_CLASS.items()
+}
+
+
+@tf_export(v1=["errors.exception_type_from_error_code"])
+def exception_type_from_error_code(error_code):
+  return _CODE_TO_EXCEPTION_CLASS[error_code]
+
+
+@tf_export(v1=["errors.error_code_from_exception_type"])
+def error_code_from_exception_type(cls):
+  try:
+    return _EXCEPTION_CLASS_TO_CODE[cls]
+  except KeyError:
+    warnings.warn("Unknown class exception")
+    return UnknownError(None, None, "Unknown class exception", None)
+
+
+def _make_specific_exception(node_def, op, message, error_code):
+  try:
+    exc_type = exception_type_from_error_code(error_code)
+    return exc_type(node_def, op, message)
+  except KeyError:
+    warnings.warn("Unknown error code: %d" % error_code)
+    return UnknownError(node_def, op, message, error_code)
+
+
+# Named like a function for backwards compatibility with the
+# @tf_contextlib.contextmanager version, which was switched to a class to avoid
+# some object creation overhead.
+# TODO(b/77295559): expand use of TF_Status* SWIG typemap and deprecate this.
+@tf_export(v1=["errors.raise_exception_on_not_ok_status"])  # pylint: disable=invalid-name
+class raise_exception_on_not_ok_status(object):
+  """Context manager to check for C API status."""
+
+  def __enter__(self):
+    self.status = c_api_util.ScopedTFStatus()
+    return self.status.status
+
+  def __exit__(self, type_arg, value_arg, traceback_arg):
+    try:
+      if c_api.TF_GetCode(self.status.status) != 0:
+        raise _make_specific_exception(
+            None, None, compat.as_text(c_api.TF_Message(self.status.status)),
+            c_api.TF_GetCode(self.status.status))
+    # Delete the underlying status object from memory otherwise it stays alive
+    # as there is a reference to status from this from the traceback due to
+    # raise.
+    finally:
+      del self.status
+    return False  # False values do not suppress exceptions

tensorflow/python/framework/errors_test.py

@@ -0,0 +1,223 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.errors."""
+
+import gc
+import pickle
+import warnings
+
+from tensorflow.core.lib.core import error_codes_pb2
+from tensorflow.python.framework import _errors_test_helper
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.lib.io import _pywrap_file_io
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class ErrorsTest(test.TestCase):
+
+  def _CountReferences(self, typeof):
+    """Count number of references to objects of type |typeof|."""
+    objs = gc.get_objects()
+    ref_count = 0
+    for o in objs:
+      try:
+        if isinstance(o, typeof):
+          ref_count += 1
+      # Certain versions of python keeps a weakref to deleted objects.
+      except ReferenceError:
+        pass
+    return ref_count
+
+  def testUniqueClassForEachErrorCode(self):
+    for error_code, exc_type in [
+        (errors.CANCELLED, errors_impl.CancelledError),
+        (errors.UNKNOWN, errors_impl.UnknownError),
+        (errors.INVALID_ARGUMENT, errors_impl.InvalidArgumentError),
+        (errors.DEADLINE_EXCEEDED, errors_impl.DeadlineExceededError),
+        (errors.NOT_FOUND, errors_impl.NotFoundError),
+        (errors.ALREADY_EXISTS, errors_impl.AlreadyExistsError),
+        (errors.PERMISSION_DENIED, errors_impl.PermissionDeniedError),
+        (errors.UNAUTHENTICATED, errors_impl.UnauthenticatedError),
+        (errors.RESOURCE_EXHAUSTED, errors_impl.ResourceExhaustedError),
+        (errors.FAILED_PRECONDITION, errors_impl.FailedPreconditionError),
+        (errors.ABORTED, errors_impl.AbortedError),
+        (errors.OUT_OF_RANGE, errors_impl.OutOfRangeError),
+        (errors.UNIMPLEMENTED, errors_impl.UnimplementedError),
+        (errors.INTERNAL, errors_impl.InternalError),
+        (errors.UNAVAILABLE, errors_impl.UnavailableError),
+        (errors.DATA_LOSS, errors_impl.DataLossError),
+    ]:
+      # pylint: disable=protected-access
+      self.assertTrue(
+          isinstance(
+              errors_impl._make_specific_exception(None, None, None,
+                                                   error_code), exc_type))
+      # error_code_from_exception_type and exception_type_from_error_code should
+      # be consistent with operation result.
+      self.assertEqual(error_code,
+                       errors_impl.error_code_from_exception_type(exc_type))
+      # pylint: enable=protected-access
+
+  def testKnownErrorClassForEachErrorCodeInProto(self):
+    for error_code in error_codes_pb2.Code.values():
+      # pylint: disable=line-too-long
+      if error_code in (
+          error_codes_pb2.OK, error_codes_pb2.
+          DO_NOT_USE_RESERVED_FOR_FUTURE_EXPANSION_USE_DEFAULT_IN_SWITCH_INSTEAD_
+      ):
+        continue
+      # pylint: enable=line-too-long
+      with warnings.catch_warnings(record=True) as w:
+        # pylint: disable=protected-access
+        exc = errors_impl._make_specific_exception(None, None, None, error_code)
+        # pylint: enable=protected-access
+      self.assertEqual(0, len(w))  # No warning is raised.
+      self.assertTrue(isinstance(exc, errors_impl.OpError))
+      self.assertTrue(errors_impl.OpError in exc.__class__.__bases__)
+
+  def testUnknownErrorCodeCausesWarning(self):
+    with warnings.catch_warnings(record=True) as w:
+      # pylint: disable=protected-access
+      exc = errors_impl._make_specific_exception(None, None, None, 37)
+      # pylint: enable=protected-access
+    self.assertEqual(1, len(w))
+    self.assertTrue("Unknown error code: 37" in str(w[0].message))
+    self.assertTrue(isinstance(exc, errors_impl.OpError))
+
+    with warnings.catch_warnings(record=True) as w:
+      # pylint: disable=protected-access
+      exc = errors_impl.error_code_from_exception_type("Unknown")
+      # pylint: enable=protected-access
+    self.assertEqual(1, len(w))
+    self.assertTrue("Unknown class exception" in str(w[0].message))
+    self.assertTrue(isinstance(exc, errors_impl.OpError))
+
+  def testStatusDoesNotLeak(self):
+    try:
+      _pywrap_file_io.DeleteFile(compat.as_bytes("/DOES_NOT_EXIST/"))
+    except:
+      pass
+    gc.collect()
+    self.assertEqual(0, self._CountReferences(c_api_util.ScopedTFStatus))
+
+  def testPickleable(self):
+    for error_code in [
+        errors.CANCELLED,
+        errors.UNKNOWN,
+        errors.INVALID_ARGUMENT,
+        errors.DEADLINE_EXCEEDED,
+        errors.NOT_FOUND,
+        errors.ALREADY_EXISTS,
+        errors.PERMISSION_DENIED,
+        errors.UNAUTHENTICATED,
+        errors.RESOURCE_EXHAUSTED,
+        errors.FAILED_PRECONDITION,
+        errors.ABORTED,
+        errors.OUT_OF_RANGE,
+        errors.UNIMPLEMENTED,
+        errors.INTERNAL,
+        errors.UNAVAILABLE,
+        errors.DATA_LOSS,
+    ]:
+      # pylint: disable=protected-access
+      exc = errors_impl._make_specific_exception(None, None, None, error_code)
+      # pylint: enable=protected-access
+      unpickled = pickle.loads(pickle.dumps(exc))
+      self.assertEqual(exc.node_def, unpickled.node_def)
+      self.assertEqual(exc.op, unpickled.op)
+      self.assertEqual(exc.message, unpickled.message)
+      self.assertEqual(exc.error_code, unpickled.error_code)
+
+  def testErrorPayloadsFromStatus(self):
+    for code, expected_exception in [
+        (1, errors.CancelledError),
+        (2, errors.UnknownError),
+        (3, errors.InvalidArgumentError),
+        (4, errors.DeadlineExceededError),
+        (5, errors.NotFoundError),
+        (6, errors.AlreadyExistsError),
+        (7, errors.PermissionDeniedError),
+        (16, errors.UnauthenticatedError),
+        (8, errors.ResourceExhaustedError),
+        (9, errors.FailedPreconditionError),
+        (10, errors.AbortedError),
+        (11, errors.OutOfRangeError),
+        (12, errors.UnimplementedError),
+        (13, errors.InternalError),
+        (14, errors.UnavailableError),
+        (15, errors.DataLossError),
+    ]:
+      with self.assertRaises(expected_exception) as error:
+        _errors_test_helper.TestRaiseFromStatus(code)
+      self.assertEqual(error.exception.experimental_payloads[b"key1"],
+                       b"value1")
+      self.assertEqual(error.exception.experimental_payloads[b"key2"],
+                       b"value2")
+
+  def testErrorPayloadsFromTFStatus(self):
+    for code, expected_exception in [
+        (1, errors.CancelledError),
+        (2, errors.UnknownError),
+        (3, errors.InvalidArgumentError),
+        (4, errors.DeadlineExceededError),
+        (5, errors.NotFoundError),
+        (6, errors.AlreadyExistsError),
+        (7, errors.PermissionDeniedError),
+        (16, errors.UnauthenticatedError),
+        (8, errors.ResourceExhaustedError),
+        (9, errors.FailedPreconditionError),
+        (10, errors.AbortedError),
+        (11, errors.OutOfRangeError),
+        (12, errors.UnimplementedError),
+        (13, errors.InternalError),
+        (14, errors.UnavailableError),
+        (15, errors.DataLossError),
+    ]:
+      with self.assertRaises(expected_exception) as error:
+        _errors_test_helper.TestRaiseFromTFStatus(code)
+      self.assertEqual(error.exception.experimental_payloads[b"key1"],
+                       b"value1")
+      self.assertEqual(error.exception.experimental_payloads[b"key2"],
+                       b"value2")
+
+  def testErrorPayloadsDefaultValue(self):
+    for exception_type in [
+        (errors.CancelledError),
+        (errors.UnknownError),
+        (errors.InvalidArgumentError),
+        (errors.DeadlineExceededError),
+        (errors.NotFoundError),
+        (errors.AlreadyExistsError),
+        (errors.PermissionDeniedError),
+        (errors.UnauthenticatedError),
+        (errors.ResourceExhaustedError),
+        (errors.FailedPreconditionError),
+        (errors.AbortedError),
+        (errors.OutOfRangeError),
+        (errors.UnimplementedError),
+        (errors.InternalError),
+        (errors.UnavailableError),
+        (errors.DataLossError),
+    ]:
+      e = exception_type(None, None, None)
+      self.assertEqual(type(e.experimental_payloads), dict)
+      self.assertEqual(len(e.experimental_payloads), 0)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/framework/errors_test_helper.cc

@@ -0,0 +1,43 @@
+/* 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 "absl/status/status.h"
+#include "absl/strings/cord.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+namespace tensorflow {
+PYBIND11_MODULE(_errors_test_helper, m) {
+  m.def("TestRaiseFromStatus", [](int code) {
+    absl::Status status(static_cast<absl::StatusCode>(code), "test message");
+    status.SetPayload("key1", absl::Cord("value1"));
+    status.SetPayload("key2", absl::Cord("value2"));
+    MaybeRaiseRegisteredFromStatus(status);
+    return 0;
+  });
+
+  m.def("TestRaiseFromTFStatus", [](int code) {
+    TF_Status *tf_status = TF_NewStatus();
+    TF_SetStatus(tf_status, static_cast<TF_Code>(code), "test_message");
+    TF_SetPayload(tf_status, "key1", "value1");
+    TF_SetPayload(tf_status, "key2", "value2");
+    MaybeRaiseRegisteredFromTFStatus(tf_status);
+    TF_DeleteStatus(tf_status);
+    return 0;
+  });
+}
+}  // namespace tensorflow

tensorflow/python/framework/experimental/BUILD

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

tensorflow/python/framework/experimental/_math_ops.pyi

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

tensorflow/python/framework/experimental/_nn_ops.pyi

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

tensorflow/python/framework/experimental/_unified_api.pyi

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

tensorflow/python/framework/experimental/context_stack.py

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

tensorflow/python/framework/experimental/def_function.py

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

tensorflow/python/framework/experimental/gradient_registry.py

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

tensorflow/python/framework/experimental/graph_building_test.py

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

tensorflow/python/framework/experimental/math_ops.cc

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

tensorflow/python/framework/experimental/math_ops.py

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

tensorflow/python/framework/experimental/nn_ops.cc

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

tensorflow/python/framework/experimental/nn_ops.py

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

tensorflow/python/framework/experimental/tape.cc

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

tensorflow/python/framework/experimental/tape.py

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

tensorflow/python/framework/experimental/thread_local_stack.py

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

tensorflow/python/framework/experimental/unified_api.cc

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

tensorflow/python/framework/experimental/unified_api_test.py

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

tensorflow/python/framework/extension_type_field.py

@@ -0,0 +1,418 @@
+# 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.
+# ==============================================================================
+"""Metadata about fields for user-defined ExtensionType classes."""
+
+import collections
+import collections.abc
+import enum
+import typing
+
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import immutable_dict
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import type_spec
+from tensorflow.python.util import type_annotations
+
+# These names may not be used as the name for a ExtensionType field (to prevent
+# name clashes).  All names beginning with `'_tf_extension_type'` are also
+# reserved.
+RESERVED_FIELD_NAMES = [
+    'self',
+    # Name of the nested TypeSpec class.
+    'Spec',
+    # Names defined by the CompositeTensor base class.
+    '_type_spec',
+    '_shape_invariant_to_type_spec',
+    '_consumers',
+    # Names defined by the TypeSpec base class.
+    'value_type',
+    'is_compatible_with',
+    'most_specific_compatible_type',
+    '_with_tensor_ranks_only',
+    '_to_components',
+    '_from_components',
+    '_component_specs',
+    '_to_tensor_list',
+    '_from_tensor_list',
+    '_from_compatible_tensor_list',
+    '_flat_tensor_specs',
+    '_serialize',
+    '_deserialize',
+    '_to_legacy_output_types',
+    '_to_legacy_output_shapes',
+    '_to_legacy_output_classes',
+    # Used by Keras
+    '_keras_mask'
+]
+
+
+class Sentinel(object):
+  """Sentinel value that's not equal (w/ `is`) to any user value."""
+
+  def __init__(self, name):
+    self._name = name
+
+  def __repr__(self):
+    return self._name
+
+
+_NoneType = type(None)
+
+
+def _issubclass(cls, clsinfo):
+  """Internal issubclass that doesn't raise TypeError."""
+  try:
+    return issubclass(cls, clsinfo)
+  except TypeError:
+    # issubclass with GenericAlias instances raises TypeError. For example,
+    # `issubclass(tuple[int], composite_tensor.CompositeTensor)`.
+    return False
+
+
+# ==============================================================================
+# ExtensionTypeField
+# ==============================================================================
+class ExtensionTypeField(
+    collections.namedtuple('ExtensionTypeField',
+                           ['name', 'value_type', 'default'])):
+  """Metadata about a single field in a `tf.ExtensionType` object."""
+
+  NO_DEFAULT = Sentinel('ExtensionTypeField.NO_DEFAULT')
+
+  def __new__(cls, name, value_type, default=NO_DEFAULT):
+    """Constructs a new ExtensionTypeField containing metadata for a single field.
+
+    Args:
+      name: The name of the new field (`str`).  May not be a reserved name.
+      value_type: A python type expression constraining what values this field
+        can take.
+      default: The default value for the new field, or `NO_DEFAULT` if this
+        field has no default value.
+
+    Returns:
+      A new `ExtensionTypeField`.
+
+    Raises:
+      TypeError: If the type described by `value_type` is not currently
+          supported by `tf.ExtensionType`.
+      TypeError: If `default` is specified and its type does not match
+        `value_type`.
+    """
+    try:
+      validate_field_value_type(value_type, allow_forward_references=True)
+    except TypeError as e:
+      raise TypeError(f'In field {name!r}: {e}') from e
+
+    if default is not cls.NO_DEFAULT:
+      default = _convert_value(default, value_type,
+                               (f'default value for {name}',),
+                               _ConversionContext.DEFAULT)
+    return super(ExtensionTypeField, cls).__new__(cls, name, value_type,
+                                                  default)
+
+  @staticmethod
+  def is_reserved_name(name):
+    """Returns true if `name` is a reserved name."""
+    return name in RESERVED_FIELD_NAMES or name.lower().startswith(
+        '_tf_extension_type')
+
+
+def validate_field_value_type(value_type,
+                              in_mapping_key=False,
+                              allow_forward_references=False):
+  """Checks that `value_type` contains only supported type annotations.
+
+  Args:
+    value_type: The type annotation to check.
+    in_mapping_key: True if `value_type` is nested in the key of a mapping.
+    allow_forward_references: If false, then raise an exception if a
+      `value_type` contains a forward reference (i.e., a string literal).
+
+  Raises:
+    TypeError: If `value_type` contains an unsupported type annotation.
+  """
+  if isinstance(value_type, str) or type_annotations.is_forward_ref(value_type):
+    if allow_forward_references:
+      return
+    else:
+      raise TypeError(f'Unresolved forward reference {value_type!r}')
+
+  if value_type in (int, float, str, bytes, bool, None, _NoneType,
+                    dtypes.DType):
+    return
+  elif (value_type in (tensor.Tensor, tensor_shape.TensorShape) or
+        (isinstance(value_type, type) and
+         _issubclass(value_type, composite_tensor.CompositeTensor))):
+    if in_mapping_key:
+      raise TypeError(f'Mapping had a key {value_type.__name__!r} with type '
+                      f'{type(value_type).__name__!r}')
+  elif (type_annotations.is_generic_tuple(value_type) or
+        type_annotations.is_generic_union(value_type)):
+    type_args = type_annotations.get_generic_type_args(value_type)
+    if (len(type_args) == 2 and type_args[1] is Ellipsis and
+        type_annotations.is_generic_tuple(value_type)):  # `Tuple[X, ...]`
+      validate_field_value_type(type_args[0], in_mapping_key,
+                                allow_forward_references)
+    else:
+      for arg in type_annotations.get_generic_type_args(value_type):
+        validate_field_value_type(arg, in_mapping_key, allow_forward_references)
+  elif type_annotations.is_generic_mapping(value_type):
+    key_type, value_type = type_annotations.get_generic_type_args(value_type)
+    validate_field_value_type(key_type, True, allow_forward_references)
+    validate_field_value_type(value_type, in_mapping_key,
+                              allow_forward_references)
+  elif isinstance(value_type, type):
+    raise TypeError(f'Unsupported type annotation {value_type.__name__!r}')
+  else:
+    raise TypeError(f'Unsupported type annotation {value_type!r}')
+
+
+# ==============================================================================
+# Type-checking & conversion for ExtensionTypeField values
+# ==============================================================================
+
+
+class _ConversionContext(enum.Enum):
+  """Enum to indicate what kind of value is being converted.
+
+  Used by `_convert_fields` and `_convert_value` and their helper methods.
+  """
+  VALUE = 1  # Converting an ExtensionType field
+  SPEC = 2  # Converting an ExtensionType.Spec field
+  DEFAULT = 3  # Converting a default value for __init__
+
+
+def convert_fields(fields, field_values):
+  """Type-checks and converts each field in `field_values` (in place).
+
+  Args:
+    fields: A list of `ExtensionTypeField` objects.
+    field_values: A `dict` mapping field names to values.  Must contain an entry
+      for each field.  I.e., `set(field_values.keys())` must be equal to
+      `set([f.name for f in fields])`.
+
+  Raises:
+    ValueError: If the keys of `field_values` do not match the names of
+      the fields in `fields`.
+    TypeError: If any value in `field_values` does not have the type indicated
+      by the corresponding `ExtensionTypeField` object.
+  """
+  _convert_fields(fields, field_values, context=_ConversionContext.VALUE)
+
+
+def convert_fields_for_spec(fields, field_values):
+  """Type-checks and converts field values for a TypeSpec (in place).
+
+  This is similar to `convert_fields`, except that we expect a `TypeSpec` for
+  tensor-like types.  In particular, if the `value_type` of a field is
+  `tf.Tensor` or a `CompositeTensor` subclass, then the corresponding value in
+  `fields` is expected to contain a `TypeSpec` (rather than a value described by
+  that `TypeSpec`).
+
+  Args:
+    fields: A list of `ExtensionTypeField` objects.
+    field_values: A `dict` mapping field names to values.  Must contain an entry
+      for each field.  I.e., `set(field_values.keys())` must be equal to
+      `set([f.name for f in fields])`.
+
+  Raises:
+    ValueError: If the keys of `field_values` do not match the names of
+      the fields in `fields`.
+    TypeError: If any value in `field_values` does not have the type indicated
+      by the corresponding `ExtensionTypeField` object.
+  """
+  _convert_fields(fields, field_values, context=_ConversionContext.SPEC)
+
+
+def _convert_fields(fields, field_values, context):
+  """Type-checks and converts each field in `field_values` (in place).
+
+  Args:
+    fields: A list of `ExtensionTypeField` objects.
+    field_values: A `dict` mapping field names to values.  Must contain an entry
+      for each field.  I.e., `set(field_values.keys())` must be equal to
+      `set([f.name for f in fields])`.
+    context: _ConversionContext, indicates what kind of value we are converting.
+
+  Raises:
+    ValueError: If the keys of `field_values` do not match the names of
+      the fields in `fields`.
+    TypeError: If any value in `field_values` does not have the type indicated
+      by the corresponding `ExtensionTypeField` object.
+  """
+  converted = {}
+  if len(fields) != len(field_values):
+    _report_field_mismatches(fields, field_values)
+  for field in fields:
+    if field.name not in field_values:
+      _report_field_mismatches(fields, field_values)
+    field_value = field_values[field.name]
+    converted[field.name] = _convert_value(field_value, field.value_type,
+                                           (field.name,), context)
+  field_values.update(converted)
+
+
+def _convert_value(value, expected_type, path,
+                   context=_ConversionContext.VALUE):
+  """Type-checks and converts a value.
+
+  Args:
+    value: The value to type-check.
+    expected_type: The expected type for the value.
+    path: Tuple of `str` naming the value (used for exception messages).
+    context: _ConversionContext, indicates what kind of value we are converting.
+
+  Returns:
+    A copy of `value`, converted to the expected type.
+
+  Raises:
+    TypeError: If `value` can not be converted to the expected type.
+  """
+  assert isinstance(path, tuple)
+
+  if expected_type is None:
+    expected_type = _NoneType
+
+  if expected_type is tensor.Tensor:
+    return _convert_tensor(value, path, context)
+  elif (isinstance(expected_type, type) and
+        _issubclass(expected_type, composite_tensor.CompositeTensor)):
+    return _convert_composite_tensor(value, expected_type, path, context)
+  elif expected_type is tensor_shape.TensorShape:
+    try:
+      return tensor_shape.as_shape(value)
+    except TypeError as e:
+      raise TypeError(f"{''.join(path)}: expected 'tf.TensorShape', got "
+                      f'{type(value).__name__!r}') from e
+  elif expected_type is dtypes.DType:
+    try:
+      return dtypes.as_dtype(value)
+    except TypeError as e:
+      raise TypeError(f"{''.join(path)}: expected 'tf.DType', got "
+                      f'{type(value).__name__!r}') from e
+  elif expected_type in (int, float, bool, str, bytes, _NoneType):
+    if not isinstance(value, expected_type):
+      raise TypeError(f'{"".join(path)}: expected {expected_type.__name__!r}, '
+                      f'got {type(value).__name__!r}')
+    return value
+  elif type_annotations.is_generic_tuple(expected_type):
+    return _convert_tuple(value, expected_type, path, context)
+  elif type_annotations.is_generic_mapping(expected_type):
+    return _convert_mapping(value, expected_type, path, context)
+  elif type_annotations.is_generic_union(expected_type):
+    return _convert_union(value, expected_type, path, context)
+  else:
+    raise TypeError(f'{"".join(path)}: Unsupported type annotation '
+                    f'{expected_type!r}')
+
+
+def _convert_tensor(value, path, context):
+  """Converts `value` to a `Tensor`."""
+  if context == _ConversionContext.SPEC:
+    if not (isinstance(value, type_spec.TypeSpec) and
+            value.value_type is tensor.Tensor):
+      raise TypeError(
+          f'{"".join(path)}: expected a TensorSpec, got '
+          f'{type(value).__name__!r}')
+    return value
+
+  if not isinstance(value, tensor.Tensor):
+    if context == _ConversionContext.DEFAULT:
+      # TODO(edloper): Convert the value to a numpy array?  (Note: we can't just
+      # use `np.array(value)`, since the default dtypes for TF and numpy are
+      # different -- e.g., int->np.int64 but int->tf.int32.
+      return value
+    try:
+      value = ops.convert_to_tensor(value)
+    except (ValueError, TypeError) as e:
+      raise TypeError(f'{"".join(path)}: expected a Tensor, '
+                      f'got {type(value).__name__!r}') from e
+  return value
+
+
+def _convert_composite_tensor(value, expected_type, path, context):
+  """Converts `value` to a value of type `expected_type`."""
+  if context == _ConversionContext.SPEC:
+    if not (isinstance(value, type_spec.TypeSpec) and
+            _issubclass(value.value_type, expected_type)):
+      raise TypeError(f'{"".join(path)}: expected a TypeSpec for '
+                      f'{expected_type.__name__!r}, got '
+                      f'{type(value).__name__!r}')
+    return value
+
+  if not isinstance(value, expected_type):
+    raise TypeError(f'{"".join(path)}: expected {expected_type.__name__!r}, '
+                    f'got {type(value).__name__!r}')
+  return value
+
+
+def _convert_tuple(value, expected_type, path, context):
+  """Converts `value` to a tuple with type `expected_type`."""
+  if not isinstance(value, typing.Sequence):
+    raise TypeError(f'{"".join(path)}: expected tuple, got '
+                    f'{type(value).__name__!r}')
+  element_types = type_annotations.get_generic_type_args(expected_type)
+  if len(element_types) == 2 and element_types[1] is Ellipsis:
+    return tuple([
+        _convert_value(v, element_types[0], path + (f'[{i}]',), context)
+        for (i, v) in enumerate(value)
+    ])
+  else:
+    if len(value) != len(element_types):
+      raise TypeError(f'{"".join(path)}: expected tuple with length '
+                      f'{len(element_types)}, got {type(value).__name__!r})')
+    return tuple([
+        _convert_value(v, t, path + (f'[{i}]',), context)
+        for (i, (v, t)) in enumerate(zip(value, element_types))
+    ])
+
+
+def _convert_mapping(value, expected_type, path, context):
+  """Converts `value` to a mapping with type `expected_type`."""
+  if not isinstance(value, typing.Mapping):
+    raise TypeError(f'{"".join(path)}: expected mapping, got '
+                    f'{type(value).__name__!r}')
+  key_type, value_type = type_annotations.get_generic_type_args(expected_type)
+  return immutable_dict.ImmutableDict([
+      (_convert_value(k, key_type, path + ('[<key>]',), context),
+       _convert_value(v, value_type, path + (f'[{k!r}]',), context))
+      for (k, v) in value.items()
+  ])
+
+
+def _convert_union(value, expected_type, path, context):
+  """Converts `value` to a value with any of the types in `expected_type`."""
+  for type_option in type_annotations.get_generic_type_args(expected_type):
+    try:
+      return _convert_value(value, type_option, path, context)
+    except TypeError:
+      pass
+  raise TypeError(f'{"".join(path)}: expected {expected_type!r}, got '
+                  f'{type(value).__name__!r}')
+
+
+def _report_field_mismatches(fields, field_values):
+  """Raises an exception with mismatches between fields and field_values."""
+  expected = set(f.name for f in fields)
+  actual = set(field_values)
+  extra = actual - expected
+  if extra:
+    raise ValueError(f'Got unexpected fields: {extra}')
+  missing = expected - actual
+  if missing:
+    raise ValueError(f'Missing required fields: {missing}')

tensorflow/python/framework/extension_type_field_test.py

@@ -0,0 +1,359 @@
+# 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 tf.framework.extension_type_field."""
+
+import sys
+import typing
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import extension_type_field
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+if sys.version_info >= (3, 9):
+  _TUPLE = tuple
+else:
+  # Remove this branch once TF drops support for Python < 3.9.
+  _TUPLE = typing.Tuple
+
+_INT_FLOAT_UNION_PATTERN = r'(?:typing\.Union\[int, float\]|int \| float)'
+_INT_STR_UNION_PATTERN = r'(?:typing\.Union\[int, str\]|int \| str)'
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class ExtensionTypeFieldTest(test_util.TensorFlowTestCase,
+                             parameterized.TestCase):
+
+  @parameterized.parameters([
+      # Without default values:
+      ('x', int),
+      ('f', float),
+      ('t', tensor.Tensor),
+      # With default values:
+      ('x', int, 33),
+      ('y', float, 33.8),
+      ('t', tensor.Tensor, [[1, 2], [3, 4]]),
+      ('t', tensor.Tensor, lambda: constant_op.constant([[1, 2], [3, 4]])),
+      ('r', ragged_tensor.RaggedTensor,
+       lambda: ragged_factory_ops.constant([[1, 2], [3]])),
+      ('seq', typing.Tuple[typing.Union[int, float], ...], (33, 12.8, 9, 0)),
+      ('seq', typing.Tuple[typing.Union[int, float],
+                           ...], [33, 12.8, 9, 0], (33, 12.8, 9, 0)),
+      ('seq', _TUPLE[typing.Union[int, float], ...], (33, 12.8, 9, 0)),
+      ('seq', _TUPLE[typing.Union[int, float], ...], (33, 12.8, 9, 0)),
+      ('s', tensor_shape.TensorShape, [1, 2], tensor_shape.TensorShape([1, 2])),
+      ('dtype', dtypes.DType, np.int32, dtypes.int32),
+  ])
+  def testConstruction(
+      self,
+      name,
+      value_type,
+      default=extension_type_field.ExtensionTypeField.NO_DEFAULT,
+      converted_default=None):
+    if callable(default):
+      default = default()  # deferred construction (contains tensor)
+    field = extension_type_field.ExtensionTypeField(name, value_type, default)
+    if converted_default is not None:
+      default = converted_default
+    self.assertEqual(field.name, name)
+    self.assertEqual(field.value_type, value_type)
+    if isinstance(default, (tensor.Tensor, ragged_tensor.RaggedTensor)):
+      self.assertAllEqual(field.default, default)
+    else:
+      self.assertEqual(field.default, default)
+
+  @parameterized.parameters([
+      ('i', int, 8.3, "default value for i: expected 'int', got 'float'"),
+      ('f', float, 8, "default value for f: expected 'float', got 'int'"),
+      ('x', int, 'hello world',
+       "default value for x: expected 'int', got 'str'"),
+      ('seq', typing.Tuple[typing.Union[int, float], ...], [33, 12.8, 'zero'],
+       (r'default value for seq\[2\]: expected '
+        + _INT_FLOAT_UNION_PATTERN + r", got 'str'")),
+      ('seq', _TUPLE[typing.Union[int, float], ...], [33, 12.8, 'zero'],
+       (r'default value for seq\[2\]: expected '
+        + _INT_FLOAT_UNION_PATTERN + r", got 'str'")),
+      ('t', tensor.TensorSpec(None, dtypes.int32),
+       lambda: constant_op.constant(0.0),
+       'Unsupported type annotation TensorSpec.*'),
+      ('x', dict, {}, "In field 'x': Unsupported type annotation 'dict'"),
+      ('y', typing.Union[int, list], 3,
+       "In field 'y': Unsupported type annotation 'list'"),
+      ('z', typing.Mapping[tensor.Tensor, int], {},
+       "In field 'z': Mapping had a key 'Tensor' with type 'type'"),
+  ])
+  def testConstructionError(self, name, value_type, default, error):
+    if callable(default):
+      default = default()  # deferred construction (contains tensor)
+    with self.assertRaisesRegex(TypeError, error):
+      extension_type_field.ExtensionTypeField(name, value_type, default)
+
+  @parameterized.parameters([
+      (r"ExtensionTypeField\(name='i', value_type=<class 'int'>, " +
+       r'default=ExtensionTypeField\.NO_DEFAULT\)', 'i', int),
+      (r"ExtensionTypeField\(name='x', value_type=(?:typing\.Tuple|tuple)\["
+       + _INT_STR_UNION_PATTERN
+       + r', \.\.\.\], default=ExtensionTypeField\.NO_DEFAULT\)',
+       'x', typing.Tuple[typing.Union[int, str], ...]),
+      (r"ExtensionTypeField\(name='j', value_type=<class 'int'>, default=3\)",
+       'j', int, 3),
+  ])
+  def testRepr(self,
+               expected_pattern,
+               name,
+               value_type,
+               default=extension_type_field.ExtensionTypeField.NO_DEFAULT):
+    field = extension_type_field.ExtensionTypeField(name, value_type, default)
+    self.assertRegex(repr(field), expected_pattern)
+
+  @parameterized.parameters([
+      ('Spec', True),
+      ('_type_spec', True),
+      ('self', True),
+      ('x', False),
+      ('_tf_extension_type_foo_bar', True),
+  ])
+  def testIsReservedName(self, name, expected):
+    self.assertEqual(
+        extension_type_field.ExtensionTypeField.is_reserved_name(name),
+        expected)
+
+
+class ValidateFieldPyTypeTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+
+  @parameterized.parameters([
+      # Simple types
+      dict(tp=int),
+      dict(tp=float),
+      dict(tp=str),
+      dict(tp=bytes),
+      dict(tp=bool),
+      dict(tp=None),
+      dict(tp=type(None)),
+      dict(tp=dtypes.DType),
+      dict(tp=tensor_shape.TensorShape),
+      dict(tp=tensor.Tensor),
+      dict(tp='A', allow_forward_references=True),
+      # Generic types
+      dict(tp=typing.Union[int, float]),
+      dict(tp=typing.Tuple[int, ...]),
+      dict(tp=typing.Tuple[int, int]),
+      dict(tp=_TUPLE[int, ...]),
+      dict(tp=_TUPLE[int, int]),
+      dict(tp=typing.Mapping[int, int]),
+      dict(tp=typing.Mapping[str, int]),
+      dict(tp=typing.Union[int, 'A'], allow_forward_references=True),
+      dict(tp=typing.Mapping['A', int], allow_forward_references=True),
+      dict(tp=typing.Union[int, typing.Tuple[typing.Tuple[int, int], ...]]),
+      dict(tp=typing.Union[int, _TUPLE[_TUPLE[int, int], ...]]),
+      dict(tp=typing.Union[int, _TUPLE[typing.Tuple[int, int], ...]]),
+      dict(tp=typing.Union[int, typing.Tuple[_TUPLE[int, int], ...]]),
+  ])
+  def testValidPytype(self, tp, allow_forward_references=False):
+    extension_type_field.validate_field_value_type(
+        tp, allow_forward_references=allow_forward_references)
+
+  @parameterized.parameters([
+      dict(tp=dict, error="Unsupported type annotation 'dict'"),
+      dict(tp=list, error="Unsupported type annotation 'list'"),
+      dict(
+          tp=typing.Union[int, list],
+          error="Unsupported type annotation 'list'"),
+      dict(
+          tp=typing.Tuple[typing.Tuple[int, int, dict], ...],
+          error="Unsupported type annotation 'dict'"),
+      dict(
+          tp=_TUPLE[_TUPLE[int, int, dict], ...],
+          error="Unsupported type annotation 'dict'"),
+      dict(tp='A', error='Unresolved forward reference .*'),
+      dict(tp=typing.Union[int, 'A'], error='Unresolved forward reference .*'),
+      dict(tp=typing.Mapping[tensor.Tensor, int],
+           error="Mapping had a key 'Tensor' with type 'type'"),
+      dict(
+          tp=typing.Mapping[tensor_shape.TensorShape, int],
+          error="Mapping had a key 'TensorShape' with type 'ABCMeta'"),
+  ])
+  def testInvalidPytype(self, tp, error):
+    with self.assertRaisesRegex(TypeError, error):
+      extension_type_field.validate_field_value_type(tp)
+
+
+class FieldValueConverterTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+
+  @parameterized.parameters([
+      ({
+          'x': 1
+      }, "Missing required fields: {'y'}"),
+      ({
+          'x': 1,
+          'y': 2.0,
+          'z': 3
+      }, "Got unexpected fields: {'z'}"),
+  ])
+  def testConvertFieldsMismatch(self, field_values, error):
+    fields = [
+        extension_type_field.ExtensionTypeField('x', int),
+        extension_type_field.ExtensionTypeField('y', float)
+    ]
+    with self.assertRaisesRegex(ValueError, error):
+      extension_type_field.convert_fields(fields, field_values)
+
+  @parameterized.parameters([
+      (12, int),
+      (5.3, float),
+      ('foo', str),
+      (None, None),
+      (True, bool),
+      ([1, 2, 3], tensor.Tensor),
+      (lambda: constant_op.constant([1, 2, 3]), tensor.Tensor),
+      (lambda: ragged_factory_ops.constant([[1, 2], [3]]),
+       ragged_tensor.RaggedTensor),
+      ([1, 2, 3], typing.Tuple[int, ...], (1, 2, 3)),
+      ((1, 2, 3), typing.Tuple[int, int, int], (1, 2, 3)),
+      ([1, 2, 3], _TUPLE[int, ...], (1, 2, 3)),
+      ((1, 2, 3), _TUPLE[int, int, int], (1, 2, 3)),
+      ({
+          'a': 12
+      }, typing.Mapping[str, int]),
+      ({
+          'a': (12, 3.0)
+      }, typing.Mapping[str, typing.Tuple[int, float]]),
+      ({
+          'a': (12, 3.0)
+      }, typing.Mapping[str, _TUPLE[int, float]]),
+      (tensor_shape.TensorShape([1, 2]), tensor_shape.TensorShape,
+       tensor_shape.TensorShape([1, 2])),
+      ([1, 2], tensor_shape.TensorShape, tensor_shape.TensorShape([1, 2])),
+      (dtypes.int32, dtypes.DType, dtypes.int32),
+      (np.int32, dtypes.DType, dtypes.int32),
+  ])
+  def testConvertValue(self, value, value_type, expected=None):
+    if callable(value):
+      value = value()  # deferred construction (contains tensor)
+    if expected is None:
+      expected = value
+    converted = extension_type_field._convert_value(value, value_type, ('x',))
+    if isinstance(converted, (tensor.Tensor, ragged_tensor.RaggedTensor)):
+      self.assertAllEqual(converted, expected)
+    else:
+      self.assertEqual(converted, expected)
+
+  @parameterized.parameters([
+      (12, int),
+      (5.3, float),
+      ('foo', str),
+      (None, None),
+      (True, bool),
+      (tensor.TensorSpec([5]), tensor.Tensor),
+      (ragged_tensor.RaggedTensorSpec([5, None]), ragged_tensor.RaggedTensor),
+      ([1, 2, 3], typing.Tuple[int, ...], (1, 2, 3)),
+      ((1, 2, 3), typing.Tuple[int, int, int], (1, 2, 3)),
+      ([1, 2, 3], _TUPLE[int, ...], (1, 2, 3)),
+      ((1, 2, 3), _TUPLE[int, int, int], (1, 2, 3)),
+      ({
+          'a': 12
+      }, typing.Mapping[str, int]),
+      ({
+          'a': (12, 3.0)
+      }, typing.Mapping[str, typing.Tuple[int, float]]),
+      ({
+          'a': (12, 3.0)
+      }, typing.Mapping[str, _TUPLE[int, float]]),
+      (tensor_shape.TensorShape([1, 2]), tensor_shape.TensorShape,
+       tensor_shape.TensorShape([1, 2])),
+      ([1, 2], tensor_shape.TensorShape, tensor_shape.TensorShape([1, 2])),
+      (dtypes.int32, dtypes.DType, dtypes.int32),
+      (np.int32, dtypes.DType, dtypes.int32),
+  ])
+  def testConvertValueForSpec(self, value, value_type, expected=None):
+    if callable(value):
+      value = value()  # deferred construction (contains tensor)
+    if expected is None:
+      expected = value
+    converted = extension_type_field._convert_value(
+        value, value_type, ('x',),
+        extension_type_field._ConversionContext.SPEC)
+    if isinstance(converted, (tensor.Tensor, ragged_tensor.RaggedTensor)):
+      self.assertAllEqual(converted, expected)
+    else:
+      self.assertEqual(converted, expected)
+
+  @parameterized.parameters([
+      (12.3, int, "x: expected 'int', got 'float'"),
+      (12, float, "x: expected 'float', got 'int'"),
+      ([1, 2, 3.0], typing.Tuple[int, ...],
+       r"x\[2\]: expected 'int', got 'float'"),
+      ([1, 2, 3.0], _TUPLE[int, ...],
+       r"x\[2\]: expected 'int', got 'float'"),
+      ('foo', tensor_shape.TensorShape,
+       "x: expected 'tf.TensorShape', got 'str'"),
+      ('foo', dtypes.DType, "x: expected 'tf.DType', got 'str'"),
+  ])
+  def testConvertValueError(self, value, value_type, error):
+    if callable(value):
+      value = value()  # deferred construction (contains tensor)
+    with self.assertRaisesRegex(TypeError, error):
+      extension_type_field._convert_value(value, value_type, ('x',))
+
+  def testConvertFields(self):
+    fields = [
+        extension_type_field.ExtensionTypeField('x', int),
+        extension_type_field.ExtensionTypeField(
+            'y', typing.Tuple[typing.Union[int, bool], ...]),
+        extension_type_field.ExtensionTypeField(
+            'y', _TUPLE[typing.Union[int, bool], ...]),
+        extension_type_field.ExtensionTypeField('z', tensor.Tensor)
+    ]
+    field_values = {'x': 1, 'y': [1, True, 3], 'z': [[1, 2], [3, 4], [5, 6]]}
+    extension_type_field.convert_fields(fields, field_values)
+    self.assertEqual(set(field_values), set(['x', 'y', 'z']))
+    self.assertEqual(field_values['x'], 1)
+    self.assertEqual(field_values['y'], (1, True, 3))
+    self.assertIsInstance(field_values['z'], tensor.Tensor)
+    self.assertAllEqual(field_values['z'], [[1, 2], [3, 4], [5, 6]])
+
+  def testConvertFieldsForSpec(self):
+    fields = [
+        extension_type_field.ExtensionTypeField('x', int),
+        extension_type_field.ExtensionTypeField(
+            'y', typing.Tuple[typing.Union[int, bool], ...]),
+        extension_type_field.ExtensionTypeField(
+            'y', _TUPLE[typing.Union[int, bool], ...]),
+        extension_type_field.ExtensionTypeField('z', tensor.Tensor)
+    ]
+    field_values = {
+        'x': 1,
+        'y': [1, True, 3],
+        'z': tensor.TensorSpec([5, 3])
+    }
+    extension_type_field.convert_fields_for_spec(fields, field_values)
+    self.assertEqual(set(field_values), set(['x', 'y', 'z']))
+    self.assertEqual(field_values['x'], 1)
+    self.assertEqual(field_values['y'], (1, True, 3))
+    self.assertEqual(field_values['z'], tensor.TensorSpec([5, 3]))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/framework/fast_tensor_util.pyx

@@ -0,0 +1,170 @@
+# 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.
+# ==============================================================================
+
+#cython: boundscheck=False
+#cython: wraparound=False
+#cython: infer_types=True
+
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import numpy as np
+cimport numpy as np
+
+from tensorflow.python.util import compat
+
+def AppendBFloat16ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.uint16_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.half_val.append(nparray[i])
+
+
+def AppendFloat8ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.uint8_t, ndim=1] nparray):
+  tensor_proto.float8_val += nparray.tobytes()
+
+
+def AppendFloat16ArrayToTensorProto(
+    # For numpy, npy_half is a typedef for npy_uint16,
+    # see: https://github.com/numpy/numpy/blob/master/doc/source/reference/c-api.coremath.rst#half-precision-functions
+    # Because np.float16_t dosen't exist in cython, we use uint16_t here.
+    # TODO: Use np.float16_t when cython supports it.
+    tensor_proto, np.ndarray[np.uint16_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.half_val.append(nparray[i])
+
+
+def AppendFloat32ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.float32_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.float_val.append(nparray[i])
+
+
+def AppendFloat64ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.float64_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.double_val.append(nparray[i])
+
+
+def AppendInt32ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.int32_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.int_val.append(nparray[i])
+
+def AppendUInt32ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.uint32_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.uint32_val.append(nparray[i])
+
+def AppendInt64ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.int64_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.int64_val.append(nparray[i])
+
+def AppendUInt64ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.uint64_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.uint64_val.append(nparray[i])
+
+def AppendUInt8ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.uint8_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.int_val.append(nparray[i])
+
+
+def AppendUInt16ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.uint16_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.int_val.append(nparray[i])
+
+
+def AppendInt16ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.int16_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.int_val.append(nparray[i])
+
+
+def AppendInt8ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.int8_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.int_val.append(nparray[i])
+
+
+def AppendComplex64ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.complex64_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.scomplex_val.append(nparray[i].real)
+    tensor_proto.scomplex_val.append(nparray[i].imag)
+
+
+def AppendComplex128ArrayToTensorProto(
+    tensor_proto, np.ndarray[np.complex128_t, ndim=1] nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.dcomplex_val.append(nparray[i].real)
+    tensor_proto.dcomplex_val.append(nparray[i].imag)
+
+
+def AppendObjectArrayToTensorProto(tensor_proto, np.ndarray nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.string_val.append(compat.as_bytes(nparray[i]))
+
+
+def AppendBoolArrayToTensorProto(tensor_proto, nparray):
+  cdef long i, n
+  n = nparray.size
+  for i in range(n):
+    tensor_proto.bool_val.append(nparray.item(i))

tensorflow/python/framework/file_system_test.py

@@ -0,0 +1,47 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Tests for file_system."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import load_library
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class FileSystemTest(test.TestCase):
+
+  def setUp(self):
+    file_system_library = resource_loader.get_path_to_datafile(
+        "test_file_system.so")
+    load_library.load_file_system_library(file_system_library)
+
+  @test_util.run_deprecated_v1
+  def testBasic(self):
+    with self.cached_session() as sess:
+      reader = io_ops.WholeFileReader("test_reader")
+      queue = data_flow_ops.FIFOQueue(99, [dtypes.string], shapes=())
+      queue.enqueue_many([["test://foo"]]).run()
+      queue.close().run()
+      key, value = self.evaluate(reader.read(queue))
+    self.assertEqual(key, compat.as_bytes("test://foo"))
+    self.assertEqual(value, compat.as_bytes("AAAAAAAAAA"))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/framework/flexible_dtypes.py

@@ -0,0 +1,551 @@
+# 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.
+# =============================================================================
+"""Auto dtype conversion semantics for TF."""
+
+import numpy as np
+
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import weak_tensor
+from tensorflow.python.ops import variables
+from tensorflow.python.util import nest
+
+
+# PromoMode Enum that denotes safe and all mode.
+PromoMode = ops.PromoMode
+
+# Namings are similar to third_party/py/jax/_src/dtypes.py.
+_b8 = (dtypes.bool, False)
+_u8 = (dtypes.uint8, False)
+_u16 = (dtypes.uint16, False)
+_u32 = (dtypes.uint32, False)
+_u64 = (dtypes.uint64, False)
+_i8 = (dtypes.int8, False)
+_i16 = (dtypes.int16, False)
+_i32 = (dtypes.int32, False)
+_i64 = (dtypes.int64, False)
+_bf16 = (dtypes.bfloat16, False)
+_f16 = (dtypes.float16, False)
+_f32 = (dtypes.float32, False)
+_f64 = (dtypes.float64, False)
+_c64 = (dtypes.complex64, False)
+_c128 = (dtypes.complex128, False)
+# Weak dtypes
+_i32w = (dtypes.int32, True)
+_i64w = (dtypes.int64, True)
+_f32w = (dtypes.float32, True)
+_f64w = (dtypes.float64, True)
+_c128w = (dtypes.complex128, True)
+# String
+_str = (dtypes.string, False)
+
+_all_dtypes = [
+    _b8,
+    _u8,
+    _u16,
+    _u32,
+    _u64,
+    _i8,
+    _i16,
+    _i32,
+    _i64,
+    _bf16,
+    _f16,
+    _f32,
+    _f64,
+    _c64,
+    _c128,
+    _i32w,
+    _i64w,
+    _f32w,
+    _f64w,
+    _c128w,
+]
+# Python numbers
+_pi = int  # pylint: disable=invalid-name
+_pf = float  # pylint: disable=invalid-name
+_pc = complex  # pylint: disable=invalid-name
+
+# Mappings between types_pb2.DataType values and numpy.dtypes.
+_NP_TO_TF = dtypes._NP_TO_TF  # pylint: disable=protected-access
+
+# OP(arg1, arg2) => (res, weak_type, safety level)
+# If promotion mode is SAFE, results corresponding to ALL will be disallowed.
+# This map only contains one-way dtype promotion results and is used to generate
+# _BINARY_DTYPE_RES_FULL.
+_BINARY_DTYPE_RES_HALF = {
+    _b8: {
+        _b8: (_b8, PromoMode.SAFE),
+        _u8: (_u8, PromoMode.SAFE),
+        _u16: (_u16, PromoMode.SAFE),
+        _u32: (_u32, PromoMode.SAFE),
+        _u64: (_u64, PromoMode.SAFE),
+        _i8: (_i8, PromoMode.SAFE),
+        _i16: (_i16, PromoMode.SAFE),
+        _i32: (_i32, PromoMode.SAFE),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.SAFE),
+        _f16: (_f16, PromoMode.SAFE),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_i32w, PromoMode.SAFE),
+        _i64w: (_i64w, PromoMode.SAFE),
+        _f32w: (_f32w, PromoMode.SAFE),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _u8: {
+        _u8: (_u8, PromoMode.SAFE),
+        _u16: (_u16, PromoMode.SAFE),
+        _u32: (_u32, PromoMode.SAFE),
+        _u64: (_u64, PromoMode.SAFE),
+        _i8: (_i16, PromoMode.ALL),
+        _i16: (_i16, PromoMode.SAFE),
+        _i32: (_i32, PromoMode.SAFE),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.SAFE),
+        _f16: (_f16, PromoMode.SAFE),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_u8, PromoMode.SAFE),
+        _i64w: (_u8, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _u16: {
+        _u16: (_u16, PromoMode.SAFE),
+        _u32: (_u32, PromoMode.SAFE),
+        _u64: (_u64, PromoMode.SAFE),
+        _i8: (_i32, PromoMode.ALL),
+        _i16: (_i32, PromoMode.ALL),
+        _i32: (_i32, PromoMode.SAFE),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.ALL),
+        _f16: (_f16, PromoMode.ALL),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_u16, PromoMode.SAFE),
+        _i64w: (_u16, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _u32: {
+        _u32: (_u32, PromoMode.SAFE),
+        _u64: (_u64, PromoMode.SAFE),
+        _i8: (_i64, PromoMode.ALL),
+        _i16: (_i64, PromoMode.ALL),
+        _i32: (_i64, PromoMode.ALL),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.ALL),
+        _f16: (_f16, PromoMode.ALL),
+        _f32: (_f32, PromoMode.ALL),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.ALL),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_u32, PromoMode.SAFE),
+        _i64w: (_u32, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _u64: {
+        _u64: (_u64, PromoMode.SAFE),
+        _i8: (_f64w, PromoMode.ALL),
+        _i16: (_f64w, PromoMode.ALL),
+        _i32: (_f64w, PromoMode.ALL),
+        _i64: (_f64w, PromoMode.ALL),
+        _bf16: (_bf16, PromoMode.ALL),
+        _f16: (_f16, PromoMode.ALL),
+        _f32: (_f32, PromoMode.ALL),
+        _f64: (_f64, PromoMode.ALL),
+        _c64: (_c64, PromoMode.ALL),
+        _c128: (_c128, PromoMode.ALL),
+        _i32w: (_u64, PromoMode.SAFE),
+        _i64w: (_u64, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.ALL),
+        _c128w: (_c128w, PromoMode.ALL),
+    },
+    _i8: {
+        _i8: (_i8, PromoMode.SAFE),
+        _i16: (_i16, PromoMode.SAFE),
+        _i32: (_i32, PromoMode.SAFE),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.SAFE),
+        _f16: (_f16, PromoMode.SAFE),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_i8, PromoMode.SAFE),
+        _i64w: (_i8, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _i16: {
+        _i16: (_i16, PromoMode.SAFE),
+        _i32: (_i32, PromoMode.SAFE),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.ALL),
+        _f16: (_f16, PromoMode.ALL),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_i16, PromoMode.SAFE),
+        _i64w: (_i16, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _i32: {
+        _i32: (_i32, PromoMode.SAFE),
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.ALL),
+        _f16: (_f16, PromoMode.ALL),
+        _f32: (_f32, PromoMode.ALL),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.ALL),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_i32, PromoMode.SAFE),
+        _i64w: (_i32, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _i64: {
+        _i64: (_i64, PromoMode.SAFE),
+        _bf16: (_bf16, PromoMode.ALL),
+        _f16: (_f16, PromoMode.ALL),
+        _f32: (_f32, PromoMode.ALL),
+        _f64: (_f64, PromoMode.ALL),
+        _c64: (_c64, PromoMode.ALL),
+        _c128: (_c128, PromoMode.ALL),
+        _i32w: (_i64, PromoMode.SAFE),
+        _i64w: (_i64, PromoMode.SAFE),
+        _f32w: (_f64w, PromoMode.ALL),
+        _f64w: (_f64w, PromoMode.ALL),
+        _c128w: (_c128w, PromoMode.ALL),
+    },
+    _bf16: {
+        _bf16: (_bf16, PromoMode.SAFE),
+        _f16: (_f32, PromoMode.ALL),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_bf16, PromoMode.SAFE),
+        _i64w: (_bf16, PromoMode.SAFE),
+        _f32w: (_bf16, PromoMode.SAFE),
+        _f64w: (_bf16, PromoMode.SAFE),
+        _c128w: (_c64, PromoMode.ALL),
+    },
+    _f16: {
+        _f16: (_f16, PromoMode.SAFE),
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_f16, PromoMode.SAFE),
+        _i64w: (_f16, PromoMode.SAFE),
+        _f32w: (_f16, PromoMode.SAFE),
+        _f64w: (_f16, PromoMode.SAFE),
+        _c128w: (_c64, PromoMode.ALL),
+    },
+    _f32: {
+        _f32: (_f32, PromoMode.SAFE),
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_f32, PromoMode.SAFE),
+        _i64w: (_f32, PromoMode.SAFE),
+        _f32w: (_f32, PromoMode.SAFE),
+        _f64w: (_f32, PromoMode.SAFE),
+        _c128w: (_c64, PromoMode.ALL),
+    },
+    _f64: {
+        _f64: (_f64, PromoMode.SAFE),
+        _c64: (_c128, PromoMode.ALL),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_f64, PromoMode.SAFE),
+        _i64w: (_f64, PromoMode.SAFE),
+        _f32w: (_f64, PromoMode.SAFE),
+        _f64w: (_f64, PromoMode.SAFE),
+        _c128w: (_c128, PromoMode.SAFE),
+    },
+    _c64: {
+        _c64: (_c64, PromoMode.SAFE),
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_c64, PromoMode.SAFE),
+        _i64w: (_c64, PromoMode.SAFE),
+        _f32w: (_c64, PromoMode.SAFE),
+        _f64w: (_c64, PromoMode.SAFE),
+        _c128w: (_c64, PromoMode.SAFE),
+    },
+    _c128: {
+        _c128: (_c128, PromoMode.SAFE),
+        _i32w: (_c128, PromoMode.SAFE),
+        _i64w: (_c128, PromoMode.SAFE),
+        _f32w: (_c128, PromoMode.SAFE),
+        _f64w: (_c128, PromoMode.SAFE),
+        _c128w: (_c128, PromoMode.SAFE),
+    },
+    _i32w: {
+        _i32w: (_i32w, PromoMode.SAFE),
+        _i64w: (_i64w, PromoMode.SAFE),
+        _f32w: (_f32w, PromoMode.SAFE),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _i64w: {
+        _i64w: (_i64w, PromoMode.SAFE),
+        _f32w: (_f32w, PromoMode.SAFE),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _f32w: {
+        _f32w: (_f32w, PromoMode.SAFE),
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _f64w: {
+        _f64w: (_f64w, PromoMode.SAFE),
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+    _c128w: {
+        _c128w: (_c128w, PromoMode.SAFE),
+    },
+}
+
+# A full promotion table that contains two-way mappings.
+# This table can be directly used for NumPy types as well, because
+# e.g. `np.int32` == `tf.int32`.
+_BINARY_DTYPE_RES_FULL = {}
+
+
+def _initialize():
+  """Generate the rest of the promotion table from the one-way promotion table.
+
+  Returns: None
+  """
+  for dtype1 in _all_dtypes:
+    _BINARY_DTYPE_RES_FULL[dtype1] = {}
+    for dtype2 in _all_dtypes:
+      try:
+        res = _BINARY_DTYPE_RES_HALF[dtype1][dtype2]
+      except KeyError:
+        res = _BINARY_DTYPE_RES_HALF[dtype2][dtype1]
+
+      _BINARY_DTYPE_RES_FULL[dtype1][dtype2] = res
+
+  # We do not support any conversions between string and others dtypes.
+  _BINARY_DTYPE_RES_FULL[_str] = {_str: (_str, PromoMode.SAFE)}
+
+
+_initialize()
+
+# The name np.string_ and np.unicode_ were available as aliases to np.bytes_ and
+# np.str_; however, they are removed in numpy 2.0. See
+# https://numpy.org/devdocs/numpy_2_0_migration_guide.html#changes-to-namespaces.
+_all_str_dtypes = (
+    np.dtype('object_'),
+    np.dtype('bytes_'),
+    np.dtype('str_'),
+    dtypes.string,
+)
+
+
+def _is_acceptable_input_type(x):
+  """Determines if x is an acceptable input type for auto dtype conversion semantics."""
+  # List of composite types that are supported by the auto dtype conversion
+  # semantics.
+  supported_composite_types = (
+      indexed_slices.IndexedSlices,
+      weak_tensor.WeakTensor,
+      variables.Variable,
+  )
+  return isinstance(x, supported_composite_types) or not isinstance(
+      x, composite_tensor.CompositeTensor
+  )
+
+
+def _get_dtype_and_weakness(x):
+  """Returns a TF type and weak type information from x.
+
+  Args:
+    x: an input scalar, array or a NumPy/TF/Python dtype.
+
+  Raises:
+    OverflowError: if Python int x is too large to convert to int32.
+    NotImplementedError: when x is an unsupported input type.
+
+  Returns:
+    TF type and weak type information inferred from x in the form of
+    (dtype, bool).
+  """
+  if isinstance(x, weak_tensor.WeakTensor):
+    return (x.dtype, True)
+  if isinstance(x, dtypes.DType):
+    return (x, False)
+  # TODO(b/286585200): Add support for `AutoCastVariable` in Keras.
+  tf_dtype = getattr(x, 'dtype', None)
+  if isinstance(tf_dtype, dtypes.DType):
+    return (tf_dtype, False)
+  # `isinstance(tf_dtype, np.dtype)` handles classes that implement `dtype`
+  # using `np.dtype` (e.g. `xla_extension.Array`).
+  # This condition is put before e.g. python int/float because
+  # `isinstance(np.float64(1), float)` returns True.
+  if isinstance(x, (np.ndarray, np.generic)) or isinstance(tf_dtype, np.dtype):
+    # Use `dtypes.as_dtype(x.dtype)` because in `as_dtype`, the input will be
+    # compared against a list of types including TF Dtypes which are protobufs.
+    infer_dtype = dtypes.as_dtype(tf_dtype)
+    return (infer_dtype, False)
+  if isinstance(x, (bytes, str)) or tf_dtype in _all_str_dtypes:
+    return _str
+  try:
+    if x in _NP_TO_TF:
+      return (_NP_TO_TF[x], False)
+  except TypeError:
+    pass
+  # bool type check must happen before int type check because
+  # isinstance(True, int) == True (https://peps.python.org/pep-0285/).
+  if isinstance(x, bool) or x == bool:
+    return _b8
+  # TODO(b/286585058): Update implementation depending on whether Python
+  # scalars are inferred to 32 bit or 64 bit.
+  if isinstance(x, _pi):
+    if x < np.iinfo(np.int32).min or x > np.iinfo(np.int32).max:
+      raise OverflowError(f'Python int {x} too large to convert to np.int32')
+    return _i32w
+  if x == int:
+    return _i32w
+  if isinstance(x, _pf) or x == float:
+    return _f32w
+  if isinstance(x, _pc) or x == complex:
+    return _c128w
+  if isinstance(x, tensor_shape.TensorShape):
+    # Since TensorShape is always integer value, return int32.
+    return _i32
+  # Only support NumPy dtype objects with corresponding TF types.
+  if isinstance(x, np.dtype):
+    try:
+      np_dtype = dtypes.as_dtype(x)
+      return (np_dtype, False)
+    except TypeError as exc:
+      raise NotImplementedError(
+          f'Auto dtype conversion semantics does not support {x}. Try using a'
+          ' NumPy built-in dtype objects or cast them explicitly.'
+      ) from exc
+  raise NotImplementedError(
+      f'Auto dtype conversion semantics does not support {type(x)} type.'
+  )
+
+
+def _result_type_impl(*arrays_and_dtypes):
+  """Internal implementation of jnp_style_result_type.
+
+  Args:
+    *arrays_and_dtypes: A list of Tensors, Variables, NumPy arrays or python
+      numbers.
+
+  Returns:
+    The result promotion type from all the inputs.
+
+  Raises:
+    TypeError: when the promotion between the input dtypes is disabled in the
+    current mode
+
+    NotImplementedError:
+      (1) When arrays_and_dtypes contains an unsupported input type (e.g.
+      RaggedTensor).
+      (2) When there isn't a possible promotion for the input dtypes.
+  """
+  promo_safety_mode = ops.get_dtype_conversion_mode()
+  # Drop None inputs and check if input type is supported.
+  valid_arrays_and_dtypes = []
+  for inp in arrays_and_dtypes:
+    if inp is not None:
+      if _is_acceptable_input_type(inp):
+        valid_arrays_and_dtypes.append(inp)
+      else:
+        raise NotImplementedError(
+            'Auto dtype conversion semantics does not support'
+            f' {type(inp)} type.'
+        )
+
+  dtypes_and_is_weak = [
+      _get_dtype_and_weakness(x) for x in nest.flatten(valid_arrays_and_dtypes)
+  ]
+
+  # If there are no valid inputs, return f32.
+  if not dtypes_and_is_weak:
+    # If dtypes_and_is_weak is an empty list, return weakly-typed f32.
+    dtypes_and_is_weak = [(dtypes.float32, True)]
+
+  res = dtypes_and_is_weak[0]
+  for arg in dtypes_and_is_weak[1:]:
+    # Use `base_dtype` in case of `ref` types (for e.g. ref_variables).
+    res = (res[0].base_dtype, res[1])
+    arg = (arg[0].base_dtype, arg[1])
+    try:
+      res_next, allowed_mode = _BINARY_DTYPE_RES_FULL[res][arg]
+    except KeyError as exc:
+      # Throw NotImplementedError When there isn't a possible promotion for the
+      # input dtypes. We will proceed with the default system promotion if
+      # NotImplementedError is thrown.
+      raise NotImplementedError(
+          f'Implicit Conversion between {res[0]} and {arg[0]} is '
+          'not allowed. Please convert the input manually if you '
+          'need to.'
+      ) from exc
+    if allowed_mode.value > promo_safety_mode.value:
+      raise TypeError(
+          f'In promotion mode {promo_safety_mode}, implicit dtype '
+          f'promotion between ({res[0]}, weak={res[1]}) and '
+          f'({arg[0]}, weak={arg[1]}) is disallowed. '
+          'You need to explicitly specify the dtype in your op, '
+          'or relax your dtype promotion rules (such as from SAFE '
+          'mode to ALL mode).'
+      )
+    res = res_next
+
+  return res
+
+
+def result_type(*arrays_and_dtypes):
+  """Determine the result promotion dtype using the JNP-like promotion system.
+
+  Args:
+    *arrays_and_dtypes: A list of Tensors, Variables, NumPy arrays or python
+      numbers.
+
+  Returns:
+    The result promotion type from all the inputs.
+  """
+  # Make sure to catch NotImplementedError when using this method to account for
+  # inputs that are not supported yet.
+  return _result_type_impl(*arrays_and_dtypes)

tensorflow/python/framework/flexible_dtypes_test.py

@@ -0,0 +1,936 @@
+# 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.
+# =============================================================================
+"""Tests for tensorflow.python.framework.flexible_dtypes."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import extension_type
+from tensorflow.python.framework import flexible_dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import weak_tensor
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import weak_tensor_test_util
+from tensorflow.python.platform import test as tf_test
+
+
+PromoMode = ops.PromoMode
+DtypeConversionTestEnv = weak_tensor_test_util.DtypeConversionTestEnv
+
+_ALL_INPUT_TYPES = [
+    constant_op.constant(True, dtype=dtypes.bool),
+    constant_op.constant(1, dtype=dtypes.uint8),
+    constant_op.constant(1, dtype=dtypes.uint16),
+    constant_op.constant(1, dtype=dtypes.uint32),
+    constant_op.constant(1, dtype=dtypes.uint64),
+    constant_op.constant(1, dtype=dtypes.int8),
+    constant_op.constant(1, dtype=dtypes.int16),
+    constant_op.constant(1, dtype=dtypes.int32),
+    constant_op.constant(1, dtype=dtypes.int64),
+    constant_op.constant(1, dtype=dtypes.bfloat16),
+    constant_op.constant(1, dtype=dtypes.float16),
+    constant_op.constant(1, dtype=dtypes.float32),
+    constant_op.constant(1, dtype=dtypes.float64),
+    constant_op.constant(1, dtype=dtypes.complex64),
+    constant_op.constant(1, dtype=dtypes.complex128),
+    constant_op.constant(1),
+    np.array(True, dtype=np.bool_),
+    np.array(1, dtype=np.uint8),
+    np.array(1, dtype=np.uint16),
+    np.array(1, dtype=np.uint32),
+    np.array(1, dtype=np.uint64),
+    np.array(1, dtype=np.int8),
+    np.array(1, dtype=np.int16),
+    np.array(1, dtype=np.int32),
+    np.array(1, dtype=np.int64),
+    np.array(1, dtype=np.float16),
+    np.array(1, dtype=np.float32),
+    np.array(1, dtype=np.float64),
+    np.array(1, dtype=np.complex64),
+    np.array(1, dtype=np.complex128),
+    np.array(1),
+    np.uint8(1),
+    np.uint16(1),
+    np.uint32(1),
+    np.uint64(1),
+    np.int8(1),
+    np.int16(1),
+    np.int32(1),
+    np.int64(1),
+    np.float16(1),
+    np.float32(1),
+    np.float64(1),
+    np.complex64(1),
+    np.complex128(1),
+    np.int_(),
+    1,
+    1.0,
+    1.0j,
+]
+
+
+class DtypesUtilTest(tf_test.TestCase, parameterized.TestCase):
+
+  # Test all possible TF dtypes in ALL mode.
+  @parameterized.parameters(
+      (dtypes.bool, dtypes.bool, (dtypes.bool, False)),
+      (dtypes.bool, dtypes.uint8, (dtypes.uint8, False)),
+      (dtypes.bool, dtypes.uint16, (dtypes.uint16, False)),
+      (dtypes.bool, dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.bool, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.bool, dtypes.int8, (dtypes.int8, False)),
+      (dtypes.bool, dtypes.int16, (dtypes.int16, False)),
+      (dtypes.bool, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.bool, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.bool, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.bool, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.bool, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.bool, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.bool, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.bool, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.uint8, dtypes.uint8, (dtypes.uint8, False)),
+      (dtypes.uint8, dtypes.uint16, (dtypes.uint16, False)),
+      (dtypes.uint8, dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.uint8, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.uint8, dtypes.int8, (dtypes.int16, False)),
+      (dtypes.uint8, dtypes.int16, (dtypes.int16, False)),
+      (dtypes.uint8, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.uint8, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.uint8, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.uint8, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.uint8, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.uint8, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.uint8, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.uint8, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.uint16, dtypes.uint16, (dtypes.uint16, False)),
+      (dtypes.uint16, dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.uint16, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.uint16, dtypes.int8, (dtypes.int32, False)),
+      (dtypes.uint16, dtypes.int16, (dtypes.int32, False)),
+      (dtypes.uint16, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.uint16, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.uint16, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.uint16, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.uint16, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.uint16, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.uint16, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.uint16, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.uint32, dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.uint32, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.uint32, dtypes.int8, (dtypes.int64, False)),
+      (dtypes.uint32, dtypes.int16, (dtypes.int64, False)),
+      (dtypes.uint32, dtypes.int32, (dtypes.int64, False)),
+      (dtypes.uint32, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.uint32, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.uint32, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.uint32, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.uint32, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.uint32, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.uint32, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.uint64, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.uint64, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.uint64, dtypes.int16, (dtypes.float64, True)),
+      (dtypes.uint64, dtypes.int32, (dtypes.float64, True)),
+      (dtypes.uint64, dtypes.int64, (dtypes.float64, True)),
+      (dtypes.uint64, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.uint64, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.uint64, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.uint64, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.uint64, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.uint64, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.int8, dtypes.int8, (dtypes.int8, False)),
+      (dtypes.int8, dtypes.int16, (dtypes.int16, False)),
+      (dtypes.int8, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.int8, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.int8, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.int8, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.int8, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.int8, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.int8, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.int8, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.int16, dtypes.int16, (dtypes.int16, False)),
+      (dtypes.int16, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.int16, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.int16, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.int16, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.int16, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.int16, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.int16, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.int16, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.int32, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.int32, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.int32, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.int32, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.int32, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.int32, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.int32, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.int32, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.int64, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.int64, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.int64, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.int64, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.int64, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.int64, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.int64, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.bfloat16, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.bfloat16, dtypes.float16, (dtypes.float32, False)),
+      (dtypes.bfloat16, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.bfloat16, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.bfloat16, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.bfloat16, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.float16, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.float16, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.float16, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.float16, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.float16, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.float32, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.float32, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.float32, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.float32, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.float64, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.float64, dtypes.complex64, (dtypes.complex128, False)),
+      (dtypes.float64, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.complex64, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.complex64, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.complex128, dtypes.complex128, (dtypes.complex128, False)),
+  )
+  def testResultTypeTFAndTF(self, a_dtype, b_dtype, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      input_a = (
+          constant_op.constant(1, dtype=a_dtype)
+          if a_dtype != dtypes.bool
+          else constant_op.constant(True)
+      )
+      input_b = (
+          constant_op.constant(2, dtype=b_dtype)
+          if b_dtype != dtypes.bool
+          else constant_op.constant(False)
+      )
+      self.assertEqual(
+          flexible_dtypes.result_type(input_a, input_b),
+          res_dtype,
+      )
+
+  # Test NP types dtype inference.
+  @parameterized.parameters(
+      (np.bool_, np.uint8, dtypes.uint8),
+      (np.uint8, np.int8, dtypes.int16),
+      (np.uint16, np.int8, dtypes.int32),
+      (np.uint32, np.float16, dtypes.float16),
+      (np.uint64, np.float16, dtypes.float16),
+      (np.uint64, np.complex64, dtypes.complex64),
+      (np.int8, np.float16, dtypes.float16),
+      (np.int16, np.float16, dtypes.float16),
+      (np.int32, np.complex64, dtypes.complex64),
+      (np.int64, np.float16, dtypes.float16),
+      (np.float16, np.complex64, dtypes.complex64),
+      (np.float32, np.float64, dtypes.float64),
+      (np.float64, np.complex128, dtypes.complex128),
+      (np.complex64, np.complex64, dtypes.complex64),
+      (np.complex64, np.complex128, dtypes.complex128),
+  )
+  def testResultTypeNPAndNP(self, a_dtype, b_dtype, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      self.assertEqual(
+          flexible_dtypes.result_type(
+              np.array(1, dtype=a_dtype), np.array(2, dtype=b_dtype)
+          ),
+          (res_dtype, False),
+      )
+
+  # Test np.array with default types.
+  # np.array(int) => i64
+  # np.array(float) => f64
+  # np.array(complex) => complex128
+  @parameterized.parameters(
+      (1, np.bool_, (dtypes.int64, False)),
+      (1, np.uint8, (dtypes.int64, False)),
+      (1, np.uint16, (dtypes.int64, False)),
+      (1, np.uint32, (dtypes.int64, False)),
+      (1, np.uint64, (dtypes.float64, True)),
+      (1, np.int8, (dtypes.int64, False)),
+      (1, np.int16, (dtypes.int64, False)),
+      (1, np.int32, (dtypes.int64, False)),
+      (1, np.int64, (dtypes.int64, False)),
+      (1, np.float16, (dtypes.float16, False)),
+      (1, np.float32, (dtypes.float32, False)),
+      (1, np.float64, (dtypes.float64, False)),
+      (1, np.complex64, (dtypes.complex64, False)),
+      (1, np.complex128, (dtypes.complex128, False)),
+      (1.0, np.bool_, (dtypes.float64, False)),
+      (1.0, np.uint8, (dtypes.float64, False)),
+      (1.0, np.uint16, (dtypes.float64, False)),
+      (1.0, np.uint32, (dtypes.float64, False)),
+      (1.0, np.uint64, (dtypes.float64, False)),
+      (1.0, np.int8, (dtypes.float64, False)),
+      (1.0, np.int16, (dtypes.float64, False)),
+      (1.0, np.int32, (dtypes.float64, False)),
+      (1.0, np.int64, (dtypes.float64, False)),
+      (1.0, np.float16, (dtypes.float64, False)),
+      (1.0, np.float32, (dtypes.float64, False)),
+      (1.0, np.float64, (dtypes.float64, False)),
+      (1.0, np.complex64, (dtypes.complex128, False)),
+      (1.0, np.complex128, (dtypes.complex128, False)),
+      (1.0j, np.bool_, (dtypes.complex128, False)),
+      (1.0j, np.uint8, (dtypes.complex128, False)),
+      (1.0j, np.uint16, (dtypes.complex128, False)),
+      (1.0j, np.uint32, (dtypes.complex128, False)),
+      (1.0j, np.uint64, (dtypes.complex128, False)),
+      (1.0j, np.int8, (dtypes.complex128, False)),
+      (1.0j, np.int16, (dtypes.complex128, False)),
+      (1.0j, np.int32, (dtypes.complex128, False)),
+      (1.0j, np.int64, (dtypes.complex128, False)),
+      (1.0j, np.float16, (dtypes.complex128, False)),
+      (1.0j, np.float32, (dtypes.complex128, False)),
+      (1.0j, np.float64, (dtypes.complex128, False)),
+      (1.0j, np.complex64, (dtypes.complex128, False)),
+      (1.0j, np.complex128, (dtypes.complex128, False)),
+  )
+  def testResultTypeNPDefaultArray(self, array_in, dtype, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      self.assertEqual(
+          flexible_dtypes.result_type(
+              np.array(array_in), np.array(1, dtype=dtype)
+          ),
+          res_dtype,
+      )
+
+  # Test Python int inputs. Note that Python int literals are converted into
+  # weak int32 type.
+  @parameterized.parameters(
+      (dtypes.bool, (dtypes.int32, True)),
+      (dtypes.uint8, (dtypes.uint8, False)),
+      (dtypes.uint16, (dtypes.uint16, False)),
+      (dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.int8, (dtypes.int8, False)),
+      (dtypes.int16, (dtypes.int16, False)),
+      (dtypes.int32, (dtypes.int32, False)),
+      (dtypes.int64, (dtypes.int64, False)),
+      (dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.float16, (dtypes.float16, False)),
+      (dtypes.float32, (dtypes.float32, False)),
+      (dtypes.float64, (dtypes.float64, False)),
+      (dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.complex128, (dtypes.complex128, False)),
+  )
+  def testResultTypePythonInt(self, input_dtype, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      t_input = (
+          constant_op.constant(2, dtype=input_dtype)
+          if input_dtype != dtypes.bool
+          else constant_op.constant(True)
+      )
+      self.assertEqual(flexible_dtypes.result_type(1, t_input), res_dtype)
+
+  # Test Python float inputs. Note that Python float literals are converted into
+  # weak float32 type.
+  @parameterized.parameters(
+      (dtypes.bool, (dtypes.float32, True)),
+      (dtypes.uint8, (dtypes.float64, True)),
+      (dtypes.uint16, (dtypes.float64, True)),
+      (dtypes.uint32, (dtypes.float64, True)),
+      (dtypes.uint64, (dtypes.float64, True)),
+      (dtypes.int8, (dtypes.float64, True)),
+      (dtypes.int16, (dtypes.float64, True)),
+      (dtypes.int32, (dtypes.float64, True)),
+      (dtypes.int64, (dtypes.float64, True)),
+      (dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.float16, (dtypes.float16, False)),
+      (dtypes.float32, (dtypes.float32, False)),
+      (dtypes.float64, (dtypes.float64, False)),
+      (dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.complex128, (dtypes.complex128, False)),
+  )
+  def testResultTypePythonFloat(self, input_dtype, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      t_input = (
+          constant_op.constant(2, dtype=input_dtype)
+          if input_dtype != dtypes.bool
+          else constant_op.constant(True)
+      )
+      self.assertEqual(flexible_dtypes.result_type(1.0, t_input), res_dtype)
+
+  # Test Python complex inputs. Note that Python complex literals are converted
+  # into weak complex128 type.
+  @parameterized.parameters(
+      (dtypes.bool, (dtypes.complex128, True)),
+      (dtypes.uint8, (dtypes.complex128, True)),
+      (dtypes.uint16, (dtypes.complex128, True)),
+      (dtypes.uint32, (dtypes.complex128, True)),
+      (dtypes.uint64, (dtypes.complex128, True)),
+      (dtypes.int8, (dtypes.complex128, True)),
+      (dtypes.int16, (dtypes.complex128, True)),
+      (dtypes.int32, (dtypes.complex128, True)),
+      (dtypes.int64, (dtypes.complex128, True)),
+      (dtypes.bfloat16, (dtypes.complex64, False)),
+      (dtypes.float16, (dtypes.complex64, False)),
+      (dtypes.float32, (dtypes.complex64, False)),
+      (dtypes.float64, (dtypes.complex128, False)),
+      (dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.complex128, (dtypes.complex128, False)),
+  )
+  def testResultTypePythonComplex(self, input_dtype, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      t_input = (
+          constant_op.constant(2, dtype=input_dtype)
+          if input_dtype != dtypes.bool
+          else constant_op.constant(True)
+      )
+      self.assertEqual(flexible_dtypes.result_type(1.0j, t_input), res_dtype)
+
+  # Test every possible weak type + TF dtype.
+  @parameterized.parameters(
+      (dtypes.int32, dtypes.bool, (dtypes.int32, True)),
+      (dtypes.int32, dtypes.uint8, (dtypes.uint8, False)),
+      (dtypes.int32, dtypes.uint16, (dtypes.uint16, False)),
+      (dtypes.int32, dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.int32, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.int32, dtypes.int8, (dtypes.int8, False)),
+      (dtypes.int32, dtypes.int16, (dtypes.int16, False)),
+      (dtypes.int32, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.int32, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.int32, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.int32, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.int32, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.int32, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.int32, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.int32, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.int64, dtypes.bool, (dtypes.int64, True)),
+      (dtypes.int64, dtypes.uint8, (dtypes.uint8, False)),
+      (dtypes.int64, dtypes.uint16, (dtypes.uint16, False)),
+      (dtypes.int64, dtypes.uint32, (dtypes.uint32, False)),
+      (dtypes.int64, dtypes.uint64, (dtypes.uint64, False)),
+      (dtypes.int64, dtypes.int8, (dtypes.int8, False)),
+      (dtypes.int64, dtypes.int16, (dtypes.int16, False)),
+      (dtypes.int64, dtypes.int32, (dtypes.int32, False)),
+      (dtypes.int64, dtypes.int64, (dtypes.int64, False)),
+      (dtypes.int32, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.int64, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.int64, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.int64, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.int64, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.int64, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.float32, dtypes.bool, (dtypes.float32, True)),
+      (dtypes.float32, dtypes.uint8, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.uint16, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.uint32, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.uint64, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.int8, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.int16, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.int32, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.int64, (dtypes.float64, True)),
+      (dtypes.float32, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.float32, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.float32, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.float32, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.float32, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.float32, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.float64, dtypes.bool, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.uint8, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.uint16, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.uint32, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.uint64, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.int8, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.int16, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.int32, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.int64, (dtypes.float64, True)),
+      (dtypes.float64, dtypes.bfloat16, (dtypes.bfloat16, False)),
+      (dtypes.float64, dtypes.float16, (dtypes.float16, False)),
+      (dtypes.float64, dtypes.float32, (dtypes.float32, False)),
+      (dtypes.float64, dtypes.float64, (dtypes.float64, False)),
+      (dtypes.float64, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.float64, dtypes.complex128, (dtypes.complex128, False)),
+      (dtypes.complex128, dtypes.bool, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.uint8, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.uint16, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.uint32, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.uint64, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.int8, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.int16, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.int32, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.int64, (dtypes.complex128, True)),
+      (dtypes.complex128, dtypes.bfloat16, (dtypes.complex64, False)),
+      (dtypes.complex128, dtypes.float16, (dtypes.complex64, False)),
+      (dtypes.complex128, dtypes.float32, (dtypes.complex64, False)),
+      (dtypes.complex128, dtypes.float64, (dtypes.complex128, False)),
+      (dtypes.complex128, dtypes.complex64, (dtypes.complex64, False)),
+      (dtypes.complex128, dtypes.complex128, (dtypes.complex128, False)),
+  )
+  def testResultTypeWeakTypesWithTF(self, weak_dtype_a, dtype_b, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      input_a = (
+          constant_op.constant(1, dtype=weak_dtype_a)
+          if weak_dtype_a != dtypes.bool
+          else constant_op.constant(True)
+      )
+      input_b = (
+          constant_op.constant(2, dtype=dtype_b)
+          if dtype_b != dtypes.bool
+          else constant_op.constant(True)
+      )
+      weak_input_a = weak_tensor.WeakTensor(input_a)
+      self.assertEqual(
+          flexible_dtypes.result_type(weak_input_a, input_b), res_dtype
+      )
+
+  # Test all the possible weak types + weak types.
+  @parameterized.parameters(
+      (dtypes.int32, dtypes.int32, dtypes.int32),
+      (dtypes.int32, dtypes.int64, dtypes.int64),
+      (dtypes.int32, dtypes.float32, dtypes.float32),
+      (dtypes.int32, dtypes.float64, dtypes.float64),
+      (dtypes.int32, dtypes.complex128, dtypes.complex128),
+      (dtypes.int64, dtypes.int32, dtypes.int64),
+      (dtypes.int64, dtypes.int64, dtypes.int64),
+      (dtypes.int64, dtypes.float32, dtypes.float32),
+      (dtypes.int64, dtypes.float64, dtypes.float64),
+      (dtypes.int64, dtypes.complex128, dtypes.complex128),
+      (dtypes.float32, dtypes.int32, dtypes.float32),
+      (dtypes.float32, dtypes.int64, dtypes.float32),
+      (dtypes.float32, dtypes.float32, dtypes.float32),
+      (dtypes.float32, dtypes.float64, dtypes.float64),
+      (dtypes.float32, dtypes.complex128, dtypes.complex128),
+      (dtypes.float64, dtypes.int32, dtypes.float64),
+      (dtypes.float64, dtypes.int64, dtypes.float64),
+      (dtypes.float64, dtypes.float32, dtypes.float64),
+      (dtypes.float64, dtypes.float64, dtypes.float64),
+      (dtypes.float64, dtypes.complex128, dtypes.complex128),
+      (dtypes.complex128, dtypes.int32, dtypes.complex128),
+      (dtypes.complex128, dtypes.int64, dtypes.complex128),
+      (dtypes.complex128, dtypes.float32, dtypes.complex128),
+      (dtypes.complex128, dtypes.float64, dtypes.complex128),
+      (dtypes.complex128, dtypes.complex128, dtypes.complex128),
+  )
+  def testResultTypeWeakTypesWithWeakTypes(self, dtype_a, dtype_b, res_dtype):
+    with DtypeConversionTestEnv('all'):
+      input_a = (
+          constant_op.constant(1, dtype=dtype_a)
+          if dtype_a != dtypes.bool
+          else constant_op.constant(True)
+      )
+      input_b = (
+          constant_op.constant(2, dtype=dtype_b)
+          if dtype_b != dtypes.bool
+          else constant_op.constant(True)
+      )
+      weak_input_a = weak_tensor.WeakTensor(input_a)
+      weak_input_b = weak_tensor.WeakTensor(input_b)
+      self.assertEqual(
+          flexible_dtypes.result_type(weak_input_a, weak_input_b),
+          (res_dtype, True),
+      )
+
+  # Test unallowed promotions in SAFE mode. Make sure exceptions are thrown.
+  @parameterized.parameters(
+      ((dtypes.uint8, False), (dtypes.int8, False)),
+      ((dtypes.uint8, False), (dtypes.float32, True)),
+      ((dtypes.uint16, False), (dtypes.int8, False)),
+      ((dtypes.uint16, False), (dtypes.int16, False)),
+      ((dtypes.uint16, False), (dtypes.bfloat16, False)),
+      ((dtypes.uint16, False), (dtypes.float16, False)),
+      ((dtypes.uint16, False), (dtypes.float32, True)),
+      ((dtypes.uint32, False), (dtypes.int8, False)),
+      ((dtypes.uint32, False), (dtypes.int16, False)),
+      ((dtypes.uint32, False), (dtypes.int32, False)),
+      ((dtypes.uint32, False), (dtypes.bfloat16, False)),
+      ((dtypes.uint32, False), (dtypes.float32, False)),
+      ((dtypes.uint32, False), (dtypes.complex64, False)),
+      ((dtypes.uint32, False), (dtypes.float32, True)),
+      ((dtypes.uint64, False), (dtypes.int8, False)),
+      ((dtypes.uint64, False), (dtypes.int16, False)),
+      ((dtypes.uint64, False), (dtypes.int32, False)),
+      ((dtypes.uint64, False), (dtypes.int64, False)),
+      ((dtypes.uint64, False), (dtypes.bfloat16, False)),
+      ((dtypes.uint64, False), (dtypes.float16, False)),
+      ((dtypes.uint64, False), (dtypes.float32, False)),
+      ((dtypes.uint64, False), (dtypes.float64, False)),
+      ((dtypes.uint64, False), (dtypes.complex64, False)),
+      ((dtypes.uint64, False), (dtypes.complex128, False)),
+      ((dtypes.uint64, False), (dtypes.float32, True)),
+      ((dtypes.uint64, False), (dtypes.float64, True)),
+      ((dtypes.uint64, False), (dtypes.complex128, True)),
+      ((dtypes.int8, False), (dtypes.float32, True)),
+      ((dtypes.int16, False), (dtypes.bfloat16, False)),
+      ((dtypes.int16, False), (dtypes.float16, False)),
+      ((dtypes.int16, False), (dtypes.float32, True)),
+      ((dtypes.int32, False), (dtypes.bfloat16, False)),
+      ((dtypes.int32, False), (dtypes.float16, False)),
+      ((dtypes.int32, False), (dtypes.float32, False)),
+      ((dtypes.int32, False), (dtypes.complex64, False)),
+      ((dtypes.int32, False), (dtypes.float32, True)),
+      ((dtypes.int64, False), (dtypes.bfloat16, False)),
+      ((dtypes.int64, False), (dtypes.float16, False)),
+      ((dtypes.int64, False), (dtypes.float32, False)),
+      ((dtypes.int64, False), (dtypes.float64, False)),
+      ((dtypes.int64, False), (dtypes.complex64, False)),
+      ((dtypes.int64, False), (dtypes.complex128, False)),
+      ((dtypes.int64, False), (dtypes.float32, True)),
+      ((dtypes.int64, False), (dtypes.float64, True)),
+      ((dtypes.int64, False), (dtypes.complex128, True)),
+      ((dtypes.bfloat16, False), (dtypes.float16, False)),
+      ((dtypes.bfloat16, False), (dtypes.complex128, True)),
+  )
+  def testResultTypeSafeModeUnallowedPromo(self, a_dtype, b_dtype):
+    with DtypeConversionTestEnv('safe'):
+      # Create Tensor of input dtypes.
+      input_a = (
+          constant_op.constant(1, dtype=a_dtype[0])
+          if a_dtype[0] != dtypes.bool
+          else constant_op.constant(True)
+      )
+      input_b = (
+          constant_op.constant(2, dtype=b_dtype[0])
+          if b_dtype[0] != dtypes.bool
+          else constant_op.constant(False)
+      )
+      # Create WeakTensors if weak = True.
+      if a_dtype[1]:
+        input_a = weak_tensor.WeakTensor(input_a)
+      if b_dtype[1]:
+        input_b = weak_tensor.WeakTensor(input_b)
+
+      with self.assertRaises(TypeError):
+        flexible_dtypes.result_type(input_a, input_b)
+
+  # Test allowed promotions in SAFE mode. Make sure no exception is thrown.
+  @parameterized.parameters(
+      ((dtypes.bool, False), (dtypes.bool, False)),
+      ((dtypes.bool, False), (dtypes.uint8, False)),
+      ((dtypes.bool, False), (dtypes.uint16, False)),
+      ((dtypes.bool, False), (dtypes.uint32, False)),
+      ((dtypes.bool, False), (dtypes.uint64, False)),
+      ((dtypes.bool, False), (dtypes.int8, False)),
+      ((dtypes.bool, False), (dtypes.int16, False)),
+      ((dtypes.bool, False), (dtypes.int32, False)),
+      ((dtypes.bool, False), (dtypes.int64, False)),
+      ((dtypes.bool, False), (dtypes.bfloat16, False)),
+      ((dtypes.bool, False), (dtypes.float16, False)),
+      ((dtypes.bool, False), (dtypes.float32, False)),
+      ((dtypes.bool, False), (dtypes.float64, False)),
+      ((dtypes.bool, False), (dtypes.complex64, False)),
+      ((dtypes.bool, False), (dtypes.complex128, False)),
+      ((dtypes.bool, False), (dtypes.int32, True)),
+      ((dtypes.bool, False), (dtypes.int64, True)),
+      ((dtypes.bool, False), (dtypes.float32, True)),
+      ((dtypes.bool, False), (dtypes.float64, True)),
+      ((dtypes.bool, False), (dtypes.complex128, True)),
+      ((dtypes.uint8, False), (dtypes.uint8, False)),
+      ((dtypes.uint8, False), (dtypes.uint16, False)),
+      ((dtypes.uint8, False), (dtypes.uint32, False)),
+      ((dtypes.uint8, False), (dtypes.uint64, False)),
+      ((dtypes.uint8, False), (dtypes.int16, False)),
+      ((dtypes.uint8, False), (dtypes.int32, False)),
+      ((dtypes.uint8, False), (dtypes.int64, False)),
+      ((dtypes.uint8, False), (dtypes.bfloat16, False)),
+      ((dtypes.uint8, False), (dtypes.float16, False)),
+      ((dtypes.uint8, False), (dtypes.float32, False)),
+      ((dtypes.uint8, False), (dtypes.float64, False)),
+      ((dtypes.uint8, False), (dtypes.complex64, False)),
+      ((dtypes.uint8, False), (dtypes.complex128, False)),
+      ((dtypes.uint8, False), (dtypes.int32, True)),
+      ((dtypes.uint8, False), (dtypes.int64, True)),
+      ((dtypes.uint8, False), (dtypes.float64, True)),
+      ((dtypes.uint8, False), (dtypes.complex128, True)),
+      ((dtypes.uint16, False), (dtypes.uint16, False)),
+      ((dtypes.uint16, False), (dtypes.uint32, False)),
+      ((dtypes.uint16, False), (dtypes.uint64, False)),
+      ((dtypes.uint16, False), (dtypes.int32, False)),
+      ((dtypes.uint16, False), (dtypes.int64, False)),
+      ((dtypes.uint16, False), (dtypes.float32, False)),
+      ((dtypes.uint16, False), (dtypes.float64, False)),
+      ((dtypes.uint16, False), (dtypes.complex64, False)),
+      ((dtypes.uint16, False), (dtypes.complex128, False)),
+      ((dtypes.uint16, False), (dtypes.int32, True)),
+      ((dtypes.uint16, False), (dtypes.int64, True)),
+      ((dtypes.uint16, False), (dtypes.float64, True)),
+      ((dtypes.uint16, False), (dtypes.complex128, True)),
+      ((dtypes.uint32, False), (dtypes.uint32, False)),
+      ((dtypes.uint32, False), (dtypes.uint64, False)),
+      ((dtypes.uint32, False), (dtypes.int64, False)),
+      ((dtypes.uint32, False), (dtypes.float64, False)),
+      ((dtypes.uint32, False), (dtypes.complex128, False)),
+      ((dtypes.uint32, False), (dtypes.int32, True)),
+      ((dtypes.uint32, False), (dtypes.int64, True)),
+      ((dtypes.uint32, False), (dtypes.float64, True)),
+      ((dtypes.uint32, False), (dtypes.complex128, True)),
+      ((dtypes.uint64, False), (dtypes.uint64, False)),
+      ((dtypes.uint64, False), (dtypes.int32, True)),
+      ((dtypes.uint64, False), (dtypes.int64, True)),
+      ((dtypes.int8, False), (dtypes.int8, False)),
+      ((dtypes.int8, False), (dtypes.int16, False)),
+      ((dtypes.int8, False), (dtypes.int32, False)),
+      ((dtypes.int8, False), (dtypes.int64, False)),
+      ((dtypes.int8, False), (dtypes.bfloat16, False)),
+      ((dtypes.int8, False), (dtypes.float16, False)),
+      ((dtypes.int8, False), (dtypes.float32, False)),
+      ((dtypes.int8, False), (dtypes.float64, False)),
+      ((dtypes.int8, False), (dtypes.complex64, False)),
+      ((dtypes.int8, False), (dtypes.complex128, False)),
+      ((dtypes.int8, False), (dtypes.int32, True)),
+      ((dtypes.int8, False), (dtypes.int64, True)),
+      ((dtypes.int8, False), (dtypes.float64, True)),
+      ((dtypes.int8, False), (dtypes.complex128, True)),
+      ((dtypes.int16, False), (dtypes.int16, False)),
+      ((dtypes.int16, False), (dtypes.int32, False)),
+      ((dtypes.int16, False), (dtypes.int64, False)),
+      ((dtypes.int16, False), (dtypes.float32, False)),
+      ((dtypes.int16, False), (dtypes.float64, False)),
+      ((dtypes.int16, False), (dtypes.complex64, False)),
+      ((dtypes.int16, False), (dtypes.complex128, False)),
+      ((dtypes.int16, False), (dtypes.int32, True)),
+      ((dtypes.int16, False), (dtypes.int64, True)),
+      ((dtypes.int16, False), (dtypes.float64, True)),
+      ((dtypes.int16, False), (dtypes.complex128, True)),
+      ((dtypes.int32, False), (dtypes.int32, False)),
+      ((dtypes.int32, False), (dtypes.int64, False)),
+      ((dtypes.int32, False), (dtypes.float64, False)),
+      ((dtypes.int32, False), (dtypes.complex128, False)),
+      ((dtypes.int32, False), (dtypes.int32, True)),
+      ((dtypes.int32, False), (dtypes.int64, True)),
+      ((dtypes.int32, False), (dtypes.float64, True)),
+      ((dtypes.int32, False), (dtypes.complex128, True)),
+      ((dtypes.int64, False), (dtypes.int64, False)),
+      ((dtypes.int64, False), (dtypes.int32, True)),
+      ((dtypes.int64, False), (dtypes.int64, True)),
+      ((dtypes.bfloat16, False), (dtypes.bfloat16, False)),
+      ((dtypes.bfloat16, False), (dtypes.float32, False)),
+      ((dtypes.bfloat16, False), (dtypes.float64, False)),
+      ((dtypes.bfloat16, False), (dtypes.complex64, False)),
+      ((dtypes.bfloat16, False), (dtypes.complex128, False)),
+      ((dtypes.bfloat16, False), (dtypes.int32, True)),
+      ((dtypes.bfloat16, False), (dtypes.int64, True)),
+      ((dtypes.bfloat16, False), (dtypes.float32, True)),
+      ((dtypes.bfloat16, False), (dtypes.float64, True)),
+      ((dtypes.float16, False), (dtypes.float16, False)),
+      ((dtypes.float16, False), (dtypes.float32, False)),
+      ((dtypes.float16, False), (dtypes.float64, False)),
+      ((dtypes.float16, False), (dtypes.complex64, False)),
+      ((dtypes.float16, False), (dtypes.complex128, False)),
+      ((dtypes.float16, False), (dtypes.int32, True)),
+      ((dtypes.float16, False), (dtypes.int64, True)),
+      ((dtypes.float16, False), (dtypes.float32, True)),
+      ((dtypes.float16, False), (dtypes.float64, True)),
+      ((dtypes.float32, False), (dtypes.float32, False)),
+      ((dtypes.float32, False), (dtypes.float64, False)),
+      ((dtypes.float32, False), (dtypes.complex64, False)),
+      ((dtypes.float32, False), (dtypes.complex128, False)),
+      ((dtypes.float32, False), (dtypes.int32, True)),
+      ((dtypes.float32, False), (dtypes.int64, True)),
+      ((dtypes.float32, False), (dtypes.float32, True)),
+      ((dtypes.float32, False), (dtypes.float64, True)),
+      ((dtypes.float64, False), (dtypes.float64, False)),
+      ((dtypes.float64, False), (dtypes.complex128, False)),
+      ((dtypes.float64, False), (dtypes.int32, True)),
+      ((dtypes.float64, False), (dtypes.int64, True)),
+      ((dtypes.float64, False), (dtypes.float32, True)),
+      ((dtypes.float64, False), (dtypes.float64, True)),
+      ((dtypes.float64, False), (dtypes.complex128, True)),
+      ((dtypes.complex64, False), (dtypes.complex64, False)),
+      ((dtypes.complex64, False), (dtypes.complex128, False)),
+      ((dtypes.complex64, False), (dtypes.int32, True)),
+      ((dtypes.complex64, False), (dtypes.int64, True)),
+      ((dtypes.complex64, False), (dtypes.float32, True)),
+      ((dtypes.complex64, False), (dtypes.float64, True)),
+      ((dtypes.complex64, False), (dtypes.complex128, True)),
+      ((dtypes.complex128, False), (dtypes.complex128, False)),
+      ((dtypes.complex128, False), (dtypes.int32, True)),
+      ((dtypes.complex128, False), (dtypes.int64, True)),
+      ((dtypes.complex128, False), (dtypes.float32, True)),
+      ((dtypes.complex128, False), (dtypes.float64, True)),
+      ((dtypes.complex128, False), (dtypes.complex128, True)),
+      ((dtypes.int32, True), (dtypes.int32, True)),
+      ((dtypes.int32, True), (dtypes.int64, True)),
+      ((dtypes.int32, True), (dtypes.float32, True)),
+      ((dtypes.int32, True), (dtypes.float64, True)),
+      ((dtypes.int32, True), (dtypes.complex128, True)),
+      ((dtypes.int64, True), (dtypes.int64, True)),
+      ((dtypes.int64, True), (dtypes.float32, True)),
+      ((dtypes.int64, True), (dtypes.float64, True)),
+      ((dtypes.int64, True), (dtypes.complex128, True)),
+      ((dtypes.float32, True), (dtypes.float32, True)),
+      ((dtypes.float32, True), (dtypes.float64, True)),
+      ((dtypes.float32, True), (dtypes.complex128, True)),
+      ((dtypes.float64, True), (dtypes.float64, True)),
+      ((dtypes.float64, True), (dtypes.complex128, True)),
+      ((dtypes.complex128, True), (dtypes.complex128, True)),
+  )
+  def testResultTypeSafeModeAllowedPromo(self, a_dtype, b_dtype):
+    with DtypeConversionTestEnv('safe'):
+      # Create Tensor of input dtypes.
+      input_a = (
+          constant_op.constant(1, dtype=a_dtype[0])
+          if a_dtype[0] != dtypes.bool
+          else constant_op.constant(True)
+      )
+      input_b = (
+          constant_op.constant(2, dtype=b_dtype[0])
+          if b_dtype[0] != dtypes.bool
+          else constant_op.constant(False)
+      )
+      # Create WeakTensors if weak = True.
+      if a_dtype[1]:
+        input_a = weak_tensor.WeakTensor(input_a)
+      if b_dtype[1]:
+        input_b = weak_tensor.WeakTensor(input_b)
+      flexible_dtypes.result_type(input_a, input_b)
+
+  # Test Python nested structure type inference.
+  def testResultTypePythonNestedStructure(self):
+    with DtypeConversionTestEnv('all'):
+      # i32* + f32* => f32*
+      self.assertEqual(
+          flexible_dtypes.result_type([1], [1.0]),
+          (dtypes.float32, True),
+      )
+      # f32* + c128* => c128*
+      self.assertEqual(
+          flexible_dtypes.result_type([1, 2.0], [1.0j]),
+          (dtypes.complex128, True),
+      )
+      self.assertEqual(
+          flexible_dtypes.result_type([[1, 1.0], [1.0, 1.0]], [1.0j]),
+          (dtypes.complex128, True),
+      )
+
+  # Test tf.variable type inference.
+  def testResultTypeVariable(self):
+    with DtypeConversionTestEnv('all'):
+      v = variables.Variable(1.0, dtype=dtypes.float32)
+      t = constant_op.constant(1, dtype=dtypes.float64)
+      self.assertEqual(
+          flexible_dtypes.result_type(v, t),
+          (dtypes.float64, False),
+      )
+
+  # Test TF Dtypes type inference.
+  def testResultTypeTFDtype(self):
+    with DtypeConversionTestEnv('all'):
+      d1 = dtypes.float32
+      d2 = dtypes.float16
+      self.assertEqual(
+          flexible_dtypes.result_type(d1, d2),
+          (dtypes.float32, False),
+      )
+
+  # Test NP dtype class type inference.
+  def testResultTypeNPDtype(self):
+    with DtypeConversionTestEnv('all'):
+      d = np.dtype(np.float32)
+      self.assertEqual(
+          flexible_dtypes.result_type(d),
+          (dtypes.float32, False),
+      )
+
+      d = np.dtype([('f1', np.int16)])
+      with self.assertRaises(NotImplementedError):
+        _ = flexible_dtypes.result_type(d)
+
+      d = np.dtype([('a', 'f8'), ('b', 'S10')])
+      with self.assertRaises(NotImplementedError):
+        _ = flexible_dtypes.result_type(d)
+
+  # Test bool type inference.
+  def testResultTypeBool(self):
+    with DtypeConversionTestEnv('all'):
+      self.assertEqual(
+          flexible_dtypes.result_type(True, False),
+          (dtypes.bool, False),
+      )
+
+  # Test Tensor shape type inference.
+  def testResultTypeTensorShape(self):
+    with DtypeConversionTestEnv('all'):
+      t = constant_op.constant([1, 2], dtype=dtypes.float64)
+      self.assertEqual(
+          flexible_dtypes.result_type(t.shape), (dtypes.int32, False)
+      )
+
+  # Test string types.
+  def testResultTypeStr(self):
+    with DtypeConversionTestEnv('all'):
+      res = flexible_dtypes.result_type('foo', 'bar')
+      self.assertEqual(res[0], dtypes.string)
+      with self.assertRaisesRegex(
+          NotImplementedError,
+          "Implicit Conversion between <dtype: 'string'> and <dtype: 'int32'>"
+          ' is not allowed. Please convert the input manually if you need to.',
+      ):
+        flexible_dtypes.result_type('foo', 1)
+
+  # Test byte types.
+  def testResultTypeBytes(self):
+    with DtypeConversionTestEnv('all'):
+      res = flexible_dtypes.result_type(b'foo', b'bar')
+      self.assertEqual(res[0], dtypes.string)
+      with self.assertRaisesRegex(
+          NotImplementedError,
+          "Implicit Conversion between <dtype: 'string'> and <dtype: 'int32'>"
+          ' is not allowed. Please convert the input manually if you need to.',
+      ):
+        flexible_dtypes.result_type(b'foo', 1)
+
+  # Test empty input.
+  def testResultTypeEmptyInput(self):
+    with DtypeConversionTestEnv('all'):
+      dtype, is_weak = flexible_dtypes.result_type()
+      self.assertEqual(dtype, dtypes.float32)
+      self.assertTrue(is_weak)
+
+  def testResultTypeUnsupportedInputType(self):
+    class MyTensor(extension_type.ExtensionType):
+      value: tensor.Tensor
+
+    with DtypeConversionTestEnv('all'):
+      a = MyTensor(constant_op.constant(1))
+      with self.assertRaisesRegex(
+          NotImplementedError,
+          f'Auto dtype conversion semantics does not support {type(a)} type.',
+      ):
+        _ = flexible_dtypes.result_type(a)
+
+  # Test v1 + v2 = v2 + v1.
+  def testCommunicativity(self):
+    with DtypeConversionTestEnv('all'):
+      for v1 in _ALL_INPUT_TYPES:
+        for v2 in _ALL_INPUT_TYPES:
+          self.assertEqual(
+              flexible_dtypes.result_type(v1, v2),
+              flexible_dtypes.result_type(v2, v1),
+          )
+
+  # Test (v1 + v2) + v3 = v1 + (v2 + v3).
+  def testAssociativity(self):
+    with DtypeConversionTestEnv('all'):
+      for v1 in _ALL_INPUT_TYPES:
+        for v2 in _ALL_INPUT_TYPES:
+          for v3 in _ALL_INPUT_TYPES:
+            all_res = [
+                flexible_dtypes.result_type(v1, v2, v3),
+                flexible_dtypes.result_type(v1, v3, v2),
+                flexible_dtypes.result_type(v2, v1, v3),
+                flexible_dtypes.result_type(v2, v3, v1),
+                flexible_dtypes.result_type(v3, v1, v2),
+                flexible_dtypes.result_type(v3, v2, v1),
+            ]
+            self.assertAllEqual(all_res[:-1], all_res[1:])
+
+
+if __name__ == '__main__':
+  tf_test.main()
+  ops.enable_eager_execution()

tensorflow/python/framework/framework_lib.py

@@ -0,0 +1,70 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# pylint: disable=unused-import,g-bad-import-order
+"""Classes and functions for building TensorFlow graphs."""
+
+# Classes used when building a Graph.
+from tensorflow.python.framework.device import DeviceSpec
+from tensorflow.python.framework.indexed_slices import IndexedSlices
+from tensorflow.python.framework.ops import Graph
+from tensorflow.python.framework.ops import Operation
+from tensorflow.python.framework.tensor import Tensor
+
+from tensorflow.python.framework.sparse_tensor import SparseTensor
+from tensorflow.python.framework.sparse_tensor import SparseTensorValue
+
+# Utilities used when building a Graph.
+from tensorflow.python.framework.indexed_slices import convert_to_tensor_or_indexed_slices
+from tensorflow.python.framework.ops import device
+from tensorflow.python.framework.ops import container
+from tensorflow.python.framework.ops import name_scope
+from tensorflow.python.framework.ops import op_scope
+from tensorflow.python.framework.ops import colocate_with
+from tensorflow.python.framework.ops import control_dependencies
+from tensorflow.python.framework.ops import get_default_graph
+from tensorflow.python.framework.ops import reset_default_graph
+from tensorflow.python.framework.ops import GraphKeys
+from tensorflow.python.framework.ops import add_to_collection
+from tensorflow.python.framework.ops import add_to_collections
+from tensorflow.python.framework.ops import get_collection
+from tensorflow.python.framework.ops import get_collection_ref
+from tensorflow.python.framework.ops import convert_to_tensor
+from tensorflow.python.framework.random_seed import get_seed
+from tensorflow.python.framework.random_seed import set_random_seed
+from tensorflow.python.framework.sparse_tensor import convert_to_tensor_or_sparse_tensor
+from tensorflow.python.framework.importer import import_graph_def
+
+# Utilities for working with Tensors
+from tensorflow.python.framework.tensor_util import make_tensor_proto
+from tensorflow.python.framework.tensor_util import MakeNdarray as make_ndarray
+
+# Needed when you defined a new Op in C++.
+from tensorflow.python.framework.ops import RegisterGradient
+from tensorflow.python.framework.ops import NotDifferentiable
+from tensorflow.python.framework.ops import NoGradient
+from tensorflow.python.framework.tensor_shape import Dimension
+from tensorflow.python.framework.tensor_shape import TensorShape
+
+# Needed when interfacing tensorflow to new array libraries
+from tensorflow.python.framework.tensor_conversion_registry import register_tensor_conversion_function
+
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.framework.dtypes import *  # pylint: disable=redefined-builtin
+
+# Load a TensorFlow plugin
+from tensorflow.python.framework.load_library import *
+# pylint: enable=wildcard-import

tensorflow/python/framework/function_def_to_graph.py

@@ -0,0 +1,409 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Utility to convert FunctionDef to GraphDef and Graph."""
+
+import itertools
+
+
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.framework import tensor_shape_pb2
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.framework import versions_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import cpp_shape_inference_pb2
+from tensorflow.python.framework import importer
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import versions
+from tensorflow.python.framework.func_graph import FuncGraph
+from tensorflow.python.ops import resource_variable_ops
+
+
+def function_def_to_graph(
+    fdef,
+    structured_input_signature=None,
+    structured_outputs=None,
+    input_shapes=None,
+    propagate_device_spec=False,
+    include_library_functions=False,
+):
+  """Converts a FunctionDef to a FuncGraph (sub-class Graph).
+
+  The returned FuncGraph's `name`, `inputs` and `outputs` fields will be set.
+  The input tensors are represented as placeholders.
+
+  Note: `FuncGraph.inputs` and `FuncGraph.captures` are not set and may be set
+  by the caller.
+
+  Args:
+    fdef: FunctionDef.
+    structured_input_signature: Optional. The structured input signature to use
+      for initializing the FuncGraph. See the docstring for FuncGraph for more
+      information.
+    structured_outputs: Optional. The structured outputs to use for initializing
+      the FuncGraph. See the docstring for FuncGraph for more information.
+    input_shapes: Optional. A list of TensorShape objects of the shapes of
+      function inputs. Defaults to the function's "_input_shapes" attribute. If
+      specified, its length must match length of `fdef.signature.input_arg`. If
+      a shape is None, the corresponding input placeholder will have unknown
+      shape.
+    propagate_device_spec: Optional. Whether to propagate assigned device
+      information when constructing a new Graph from a FunctionDef.
+    include_library_functions: Optional. Whether to include library functions in
+      the output FuncGraph. In graph mode, the library functions will be found
+      from outer graph. In eager mode, the library functions will be found from
+      eager context.
+
+  Returns:
+    A FuncGraph.
+  """
+  func_graph = FuncGraph(fdef.signature.name,
+                         structured_input_signature=structured_input_signature,
+                         structured_outputs=structured_outputs)
+  if input_shapes is None:
+    input_shapes_attr = fdef.attr.get("_input_shapes", None)
+    if input_shapes_attr is not None:
+      raw_input_shapes = input_shapes_attr.list.shape
+
+      # Replace resource handle shapes in the inputs to disable shape inference.
+      # Setting the shape to either the variable handle shape (which is always
+      # `[]`) or the variable shape can cause shape inference issues.
+      input_shapes = []
+      for input_shape, arg_def in zip(raw_input_shapes,
+                                      fdef.signature.input_arg):
+        if arg_def.type == types_pb2.DT_RESOURCE and arg_def.handle_data:
+          input_shapes.append(None)
+        else:
+          input_shapes.append(input_shape)
+
+  graph_def, nested_to_flat_tensor_name = function_def_to_graph_def(
+      fdef, input_shapes, include_library_functions=include_library_functions
+  )
+
+  with func_graph.as_default():
+    # Add all function nodes to the graph.
+    importer.import_graph_def_for_function(
+        graph_def, name="", propagate_device_spec=propagate_device_spec)
+
+    # Initialize fields specific to FuncGraph.
+
+    # inputs
+    input_tensor_names = [
+        nested_to_flat_tensor_name[arg.name] for arg in fdef.signature.input_arg
+    ]
+    func_graph.inputs = [
+        func_graph.get_tensor_by_name(name) for name in input_tensor_names
+    ]
+
+    # outputs
+    output_tensor_names = [
+        nested_to_flat_tensor_name[fdef.ret[arg.name]]
+        for arg in fdef.signature.output_arg
+    ]
+    func_graph.outputs = [
+        func_graph.get_tensor_by_name(name) for name in output_tensor_names
+    ]
+    func_graph.control_outputs = [
+        func_graph.get_operation_by_name(fdef.control_ret[ret_name])
+        for ret_name in fdef.signature.control_output
+    ]
+
+    _set_handle_data(func_graph, fdef)
+
+    for node in graph_def.node:
+      output_shapes = node.attr.get("_output_shapes", None)
+      if output_shapes is not None:
+        op = func_graph.get_operation_by_name(node.name)
+        # _output_shapes for functions can sometimes be too long because the
+        # output-intermediates-for-gradients version of the function was
+        # substituted before saving. We'll accept that here. (See b/133666530).
+        for output_index, shape in enumerate(
+            output_shapes.list.shape[:len(op.outputs)]):
+          op.outputs[output_index].set_shape(shape)
+    output_names = {}
+    for ret_arg_def, tensor_name in zip(
+        fdef.signature.output_arg, output_tensor_names):
+      output_names[ops.tensor_id(
+          func_graph.get_tensor_by_name(tensor_name))] = (
+              ret_arg_def.name)
+    func_graph._output_names = output_names  # pylint: disable=protected-access
+  return func_graph
+
+
+def is_function(fname, graph):
+  """Checks for a function definition with `fname` in the current context."""
+  if context.executing_eagerly():
+    # Eager mode: use eager context as the single source of truth.
+    return context.context().has_function(fname)
+  else:
+    # Graph mode: use outer graphs as the single source of truth.
+    while graph is not None:
+      if graph._is_function(fname):  # pylint: disable=protected-access
+        return True
+      if hasattr(graph, "outer_graph"):
+        graph = graph.outer_graph
+      else:
+        return False
+
+
+def get_function_def(fname, graph):
+  """Gets a function definition with `fname` in the current context."""
+  if context.executing_eagerly():
+    # Eager mode: use eager context as the single source of truth.
+    if context.context().has_function(fname):
+      return context.context().get_function_def(fname)
+  else:
+    # Graph mode: use outer graphs as the single source of truth.
+    while graph is not None:
+      if graph._is_function(fname):  # pylint: disable=protected-access
+        return graph._get_function(fname).cached_definition  # pylint: disable=protected-access
+      graph = getattr(graph, "outer_graph", None)
+
+
+def copy_function_def_to_graph_def_recursively(
+    func_name, graph_def, copied_functions, default_graph=None):
+  """Recursively copies `FunctionDef`s to `GraphDef`.
+
+  It copies the outermost `FunctionDef` and all nested `FunctionDef`s to
+  `graph_def`. The `copied_function` enforces that every `FunctionDef` will be
+  copied at most once. The `FunctionDef`s will be found from `default_graph` if
+  this function was called in graph mode or from eager context if this function
+  was called in eager mode.
+
+  Args:
+    func_name: The signature name of FunctionDef to be copied to `graph_def`.
+    graph_def: The GraphDef that will contain all `FunctionDef`s in its library.
+    copied_functions: A set contains all copied function names.
+    default_graph: The `tf.Graph` where all `FunctionDef`s will be found
+      in graph mode. Not used in eager mode.
+  """
+  # Custom ops may contain a func attr with an empty fname.
+  if func_name and not is_function(func_name, default_graph):
+    raise ValueError(f"Function {func_name} was not found. Please make "
+                     "sure the FunctionDef `fdef` is correct.")
+
+  # If `copied_functions` contains `func_name`, the FunctionDef has already
+  # been added to GraphDef so we simply return here.
+  if func_name in copied_functions:
+    return
+
+  copied_functions.add(func_name)
+  func_def = get_function_def(func_name, default_graph)
+  graph_def.library.function.add().CopyFrom(func_def)
+
+  for node_def in func_def.node_def:
+    op_def = default_graph.op_def_for_type(node_def.op)
+    for attr in op_def.attr:
+      if attr.type == "func":
+        func_name = node_def.attr[attr.name].func.name
+        copy_function_def_to_graph_def_recursively(
+            func_name, graph_def, copied_functions, default_graph)
+
+      elif attr.type == "list(func)":
+        for fn in node_def.attr[attr.name].list.func:
+          func_name = fn.name
+          copy_function_def_to_graph_def_recursively(
+              func_name, graph_def, copied_functions, default_graph)
+
+
+def function_def_to_graph_def(
+    fdef, input_shapes=None, include_library_functions=False
+):
+  """Convert a FunctionDef to a GraphDef.
+
+  Steps:
+  1. Creates placeholder nodes corresponding to inputs in
+     `FunctionDef.signature.input_arg`.
+  2. Adds NodeDefs in `FunctionDef.node_def` to `GraphDef.node`.
+  3. Renames inputs of all nodes to use the convention of GraphDef instead of
+     FunctionDef. See comment on `FunctionDef.node_def` on how the tensor naming
+     in FunctionDefs is different from GraphDefs.
+
+  Args:
+    fdef: FunctionDef.
+    input_shapes: Optional. A list of TensorShape objects of the shapes of
+      function inputs. If specified, its length must match length of
+      `fdef.signature.input_arg`. If a shape is None, the corresponding input
+      placeholder will have unknown shape.
+    include_library_functions: Optional. If enabled, copy `fdef` and its
+      nested `FunctionDef`s to the library functions of the returned `GraphDef`.
+      In graph mode, the functions will be found from outer graph. In eager
+      mode, the functions will be found from eager context.
+
+  Returns:
+    A tuple of (GraphDef, dict<string, string>). The dict contains a mapping
+    from nested tensor names (in FunctionDef) to flattened names (in GraphDef).
+
+  Raises:
+    ValueError: If the length of input_shapes does not match the number of
+      input_args or if the FunctionDef is invalid.
+  """
+  graph_def = graph_pb2.GraphDef()
+  graph_def.versions.CopyFrom(
+      versions_pb2.VersionDef(
+          producer=versions.GRAPH_DEF_VERSION,
+          min_consumer=versions.GRAPH_DEF_VERSION_MIN_CONSUMER))
+
+  default_graph = ops.get_default_graph()
+
+  copied_functions = set()
+
+  if input_shapes and len(input_shapes) != len(fdef.signature.input_arg):
+    raise ValueError("Length of `input_shapes` must match the number "
+                     f"of `input_arg`s in `fdef`. Got "
+                     f"{len(input_shapes)} `input_shapes` and "
+                     f"{len(fdef.signature.input_arg)} `input_arg`s.")
+
+  # 1. Create placeholders for input nodes.
+  for i, arg_def in enumerate(fdef.signature.input_arg):
+    node_def = graph_def.node.add()
+    node_def.name = arg_def.name
+    node_def.op = "Placeholder"
+    node_def.attr["dtype"].type = arg_def.type
+    if input_shapes and input_shapes[i] is not None:
+      input_shape = input_shapes[i]
+      if not isinstance(input_shape, tensor_shape_pb2.TensorShapeProto):
+        input_shape = input_shape.as_proto()
+      node_def.attr["shape"].shape.CopyFrom(input_shape)
+    arg_attrs = fdef.arg_attr[i].attr
+    for k in arg_attrs:
+      # Only copy internal attributes. Normal attributes for nodes cannot be
+      # applied to these Placeholder nodes.
+      if k == "_output_shapes":
+        if arg_attrs[k].WhichOneof("value") == "list":
+          node_def.attr["shape"].shape.CopyFrom(arg_attrs[k].list.shape[0])
+        elif arg_attrs[k].WhichOneof("value") == "shape":
+          node_def.attr["shape"].shape.CopyFrom(arg_attrs[k].shape)
+      elif k.startswith("_"):
+        node_def.attr[k].CopyFrom(arg_attrs[k])
+
+  # 2. Copy all body NodeDefs to the GraphDef.
+  graph_def.node.extend(fdef.node_def)
+
+  # 3. Perform the renaming.
+
+  # Build the tensor name mapping then flatten the tensor names.
+  # See comment on `FunctionDef.node_def` on how the tensor naming in
+  # FunctionDefs is different from GraphDefs.
+  nested_to_flat_tensor_name = {}
+
+  for arg_def in fdef.signature.input_arg:
+    nested_to_flat_tensor_name[arg_def.name] = "{}:0".format(arg_def.name)
+    control_name = "^" + arg_def.name
+    nested_to_flat_tensor_name[control_name] = control_name
+
+  for node_def in fdef.node_def:
+    graph = default_graph
+    while True:
+      f = graph._functions.get(node_def.op, None)  # pylint: disable=protected-access
+      if f is not None or not hasattr(graph, "outer_graph"):
+        break
+      graph = graph.outer_graph
+
+    if f is not None:
+      fdef = f.cached_definition
+      op_def = fdef.signature
+      if node_def.op not in copied_functions:
+        # Since this function is referenced as an op type, we have no choice but
+        # to copy it into the GraphDef if we want downstream tools to process
+        # it.
+        graph_def.library.function.add().CopyFrom(fdef)
+        copied_functions.add(node_def.op)
+        if getattr(f, "grad_func_name", None):
+          grad_def = function_pb2.GradientDef()
+          grad_def.function_name = f.name
+          grad_def.gradient_func = f.grad_func_name
+          graph_def.library.gradient.extend([grad_def])
+    else:
+      op_def = default_graph.op_def_for_type(node_def.op)  # pylint: disable=protected-access
+
+    for attr in op_def.attr:
+      if attr.type == "func":
+        fname = node_def.attr[attr.name].func.name
+        # Custom ops may contain a func attr with an empty fname.
+        if fname and not is_function(fname, default_graph):
+          raise ValueError(f"Function {fname} was not found. Please make sure "
+                           "the FunctionDef `fdef` is correct.")
+        if include_library_functions:
+          copy_function_def_to_graph_def_recursively(
+              fname, graph_def, copied_functions, default_graph)
+
+      elif attr.type == "list(func)":
+        for fn in node_def.attr[attr.name].list.func:
+          fname = fn.name
+          # Custom ops may contain a func attr with an empty fname.
+          if fname and not is_function(fname, default_graph):
+            raise ValueError(f"Function {fname} was not found. Please make "
+                             "sure the FunctionDef `fdef` is correct.")
+          if include_library_functions:
+            copy_function_def_to_graph_def_recursively(
+                fname, graph_def, copied_functions, default_graph)
+
+    # Iterate over output_args in op_def to build the map.
+    # Index of the output tensor in the flattened list of *all* output
+    # tensors of the op.
+    flattened_index = 0
+    for arg_def in op_def.output_arg:
+      num_args = _get_num_args(arg_def, node_def)
+      for i in range(num_args):
+        # Map tensor names from "node_name:output_arg_name:index" to
+        # "node_name:flattened_index".
+        nested_name = "{}:{}:{}".format(node_def.name, arg_def.name, i)
+        flat_name = "{}:{}".format(node_def.name, flattened_index)
+        nested_to_flat_tensor_name[nested_name] = flat_name
+        flattened_index += 1
+    control_name = "^" + node_def.name
+    nested_to_flat_tensor_name[control_name] = control_name
+
+  # Update inputs of all nodes in graph.
+  for node_def in graph_def.node:
+    for i in range(len(node_def.input)):
+      node_def.input[i] = nested_to_flat_tensor_name[node_def.input[i]]
+
+  return graph_def, nested_to_flat_tensor_name
+
+
+# Based on implementation in core/framework/node_def_util.cc::ComputeArgRange.
+def _get_num_args(arg_def, node_def):
+  if arg_def.number_attr:
+    return node_def.attr[arg_def.number_attr].i
+  elif arg_def.type_list_attr:
+    return len(node_def.attr[arg_def.type_list_attr].list.type)
+  elif arg_def.type_attr or arg_def.type != types_pb2.DT_INVALID:
+    return 1
+  else:
+    raise ValueError(f"Invalid arg_def:\n\n{arg_def}. Please make sure the "
+                     "FunctionDef `fdef` is correct.")
+
+
+def _set_handle_data(func_graph, fdef):
+  """Adds handle data for resource type inputs and outputs."""
+  # The shape of the handle itself is [], while the variable shape is
+  # saved in `handle_data`. Previously, the shape of the resource handle
+  # was set to `None`. Correct both shapes here.
+  for tensor, arg_def in itertools.chain(
+      zip(func_graph.inputs, fdef.signature.input_arg),
+      zip(func_graph.outputs, fdef.signature.output_arg)):
+    if arg_def.handle_data:
+      tensor.set_shape([])
+
+      shape_and_dtype = arg_def.handle_data[0]
+      handle_data = cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
+      handle_data.is_set = True
+      handle_data.shape_and_type.append(
+          cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
+              shape=shape_and_dtype.shape, dtype=shape_and_dtype.dtype))
+      resource_variable_ops._set_handle_shapes_and_types(  # pylint: disable=protected-access
+          tensor, handle_data, True)

tensorflow/python/framework/function_def_to_graph_test.py

@@ -0,0 +1,328 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.framework.function_def_to_graph."""
+
+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 function_def_to_graph
+from tensorflow.python.framework import graph_to_function_def
+from tensorflow.python.framework import op_def_library
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class FunctionDefToGraphTest(test.TestCase):
+
+  def _build_function_def(self):
+    with ops.Graph().as_default() as g:
+      # Inputs
+      x = array_ops.placeholder(dtypes.float32, name="x")
+      y = array_ops.placeholder(dtypes.float32, name="y")
+
+      # Outputs
+      sum_squares = math_ops.add_n(
+          [math_ops.pow(x, 2), math_ops.pow(y, 2)], name="sum_squares")
+      sum_cubes = math_ops.add_n(
+          [math_ops.pow(x, 3), math_ops.pow(y, 3)], name="sum_cubes")
+    fdef = graph_to_function_def.graph_to_function_def(
+        g,
+        g.get_operations(),
+        [x, y],  # Inputs
+        [sum_squares, sum_cubes])  # Outputs.
+    fdef.signature.name = "_whats_in_a_name"
+    return fdef
+
+  @test_util.run_deprecated_v1
+  def testInputsAndOutputs(self):
+    fdef = self._build_function_def()
+    g = function_def_to_graph.function_def_to_graph(fdef)
+    self.assertEqual(g.name, "_whats_in_a_name")
+    with self.session(graph=g) as sess:
+      inputs = sess.run(g.inputs, feed_dict={"x:0": 2, "y:0": 3})
+      self.assertSequenceEqual(inputs, [2.0, 3.0])
+      outputs = sess.run(g.outputs, feed_dict={"x:0": 2, "y:0": 3})
+      self.assertSequenceEqual(outputs, [13.0, 35.0])
+
+  def testShapes(self):
+    fdef = self._build_function_def()
+
+    g = function_def_to_graph.function_def_to_graph(fdef)
+    self.assertIsNone(g.inputs[0].shape.dims)  # Unknown dims.
+    self.assertIsNone(g.inputs[1].shape.dims)  # Unknown dims.
+    self.assertIsNone(g.outputs[0].shape.dims)  # Unknown dims.
+    self.assertIsNone(g.outputs[1].shape.dims)  # Unknown dims.
+
+    g = function_def_to_graph.function_def_to_graph(
+        fdef,
+        input_shapes=[
+            tensor_shape.TensorShape([5]),
+            tensor_shape.TensorShape([5])
+        ])
+    self.assertSequenceEqual(g.inputs[0].shape.dims, [5])
+    self.assertSequenceEqual(g.inputs[1].shape.dims, [5])
+    self.assertSequenceEqual(g.outputs[0].shape.dims, [5])
+    self.assertSequenceEqual(g.outputs[1].shape.dims, [5])
+
+    g = function_def_to_graph.function_def_to_graph(
+        fdef, input_shapes=[None, tensor_shape.TensorShape([5, 7])])
+    self.assertIsNone(g.inputs[0].shape.dims)
+    self.assertSequenceEqual(g.inputs[1].shape.dims, [5, 7])
+    self.assertSequenceEqual(g.outputs[0].shape.dims, [5, 7])
+    self.assertSequenceEqual(g.outputs[1].shape.dims, [5, 7])
+
+    # Should raise a ValueError if the length of input_shapes does not match
+    # the number of input args in FunctionDef.signature.input_arg.
+    with self.assertRaises(ValueError):
+      g = function_def_to_graph.function_def_to_graph(
+          fdef, input_shapes=[tensor_shape.TensorShape([5, 7])])
+
+  def testResourceHandleInputShapes(self):
+    # Test that shape inference and validation with resource handles works as
+    # expected.
+
+    # Create a graph to generate the input and handle shape attributes in the
+    # FunctionDef.
+    with ops.Graph().as_default() as g:
+      v = variables.Variable(array_ops.ones((2, 3), dtype=dtypes.float32))
+
+      @def_function.function(
+          input_signature=[tensor_spec.TensorSpec((None, 2, 2), dtypes.int32)])
+      def lookup(inp):
+        return {
+            # gather_nd expects a nonscalar shape for `v`, otherwise raises
+            # error when doing shape inference.
+            "shape inference": array_ops.gather_nd(v, inp),
+            # Triggers output shape validation. Expected shape must be [].
+            "handle": v.handle}
+
+      lookup.get_concrete_function().add_to_graph()
+      fdef = g.as_graph_def(add_shapes=True).library.function[0]
+
+    fg = function_def_to_graph.function_def_to_graph(fdef)
+    self.assertSequenceEqual(fg.inputs[0].shape.as_list(), [None, 2, 2])
+    self.assertSequenceEqual(fg.inputs[1].shape.as_list(), [])
+
+  def testIncludeLibraryFunctions(self):
+    @def_function.function
+    def g(x):
+      return x + 1
+
+    @def_function.function
+    def f(x):
+      return g(x)
+
+    cfg = g.get_concrete_function(1.0)
+    cfg.add_to_graph()
+    gname = cfg.function_def.signature.name
+    function_def = f.get_concrete_function(1.0).function_def
+    func_graph = function_def_to_graph.function_def_to_graph(
+        function_def, include_library_functions=True)
+    graph_def = func_graph.as_graph_def()
+    self.assertLen(graph_def.library.function, 1)
+    self.assertEqual(graph_def.library.function[0].signature.name, gname)
+
+  def testCopyFunctionDefToGraphDefRecursively(self):
+    @def_function.function
+    def inner(x):
+      return x + 1
+
+    @def_function.function
+    def middle(x):
+      return inner(x) + 1
+
+    @def_function.function
+    def outer(x):
+      return middle(x) + 1
+
+    @def_function.function
+    def target_func(x):
+      return x
+
+    target_graph_def = target_func.get_concrete_function(1).graph.as_graph_def()
+
+    self.assertEmpty(target_graph_def.library.function)
+
+    concrete_outer = outer.get_concrete_function(1)
+    default_graph = ops.get_default_graph()
+    # Copy FunctionDefs in `outer` to the default graph.
+    concrete_outer.add_to_graph(default_graph)
+    outer_function_name = concrete_outer.function_def.signature.name
+    copied_functions = set()
+    # Copy all FunctionDefs from `outer` to `target_func`.
+    function_def_to_graph.copy_function_def_to_graph_def_recursively(
+        outer_function_name, target_graph_def, copied_functions, default_graph
+    )
+
+    outer_graph_def = concrete_outer.graph.as_graph_def()
+    nested_function_names = {
+        f.signature.name for f in outer_graph_def.library.function
+    }
+    expected_function_names = {outer_function_name} | nested_function_names
+
+    self.assertEqual(copied_functions, expected_function_names)
+
+    target_function_names = {
+        f.signature.name for f in target_graph_def.library.function
+    }
+
+    self.assertEqual(target_function_names, expected_function_names)
+
+
+class FunctionDefToGraphDefTest(test.TestCase):
+
+  def _build_function_def(self):
+    with ops.Graph().as_default() as g:
+      # Inputs:    x    y    z
+      #            |\   |   /
+      #            | \  |  /
+      #            |  foo_1     list_output
+      #            |   / \       /       \
+      #            | d_1 e_1  a:1        a:0
+      #            |  \   |   /           |
+      #            |   \  |  /            |
+      #            |    foo_2             |
+      #            |     / \              |
+      # Outputs:   x   d_2 e_2           a:0
+
+      x = array_ops.placeholder(dtypes.float32, name="x")
+      y = array_ops.placeholder(dtypes.int32, name="y")
+      z = array_ops.placeholder(dtypes.int32, name="z")
+
+      d_1, e_1 = op_def_library.apply_op("Foo1", name="foo_1", a=x, b=y, c=z)
+
+      list_output0, list_output1 = test_ops.list_output(
+          T=[dtypes.int32, dtypes.int32], name="list_output")
+
+      d_2, e_2 = test_ops.foo1(a=d_1, b=e_1, c=list_output1, name="foo_2")
+
+    fdef = graph_to_function_def.graph_to_function_def(
+        g,
+        g.get_operations(),
+        [x, y, z],  # Inputs
+        [x, d_2, e_2, list_output0])  # Outputs.
+
+    # Assert that the FunctionDef was correctly built.
+    assert len(fdef.node_def) == 3  # 2 Foo1 nodes and 1 ListOutput node.
+    assert fdef.node_def[0].op == "Foo1"
+    assert fdef.node_def[0].input == ["x", "y", "z"]
+    assert fdef.node_def[1].op == "ListOutput"
+    assert not fdef.node_def[1].input
+    assert fdef.node_def[2].op == "Foo1"
+    assert fdef.node_def[2].input == [
+        "foo_1:d:0", "foo_1:e:0", "list_output:a:1"
+    ]
+    return fdef
+
+  def testTensorNames(self):
+    fdef = self._build_function_def()
+    g, tensor_name_map = function_def_to_graph.function_def_to_graph_def(fdef)
+
+    # Verify that inputs of body nodes are correctly renamed.
+    # foo_1
+    self.assertSequenceEqual(g.node[3].input, ["x:0", "y:0", "z:0"])
+    # foo_2
+    self.assertSequenceEqual(g.node[5].input,
+                             ["foo_1:0", "foo_1:1", "list_output:1"])
+
+    # Verify that the `tensor_name_map` has the correct mapping.
+    self.assertDictEqual(
+        tensor_name_map, {
+            "x": "x:0",
+            "^x": "^x",
+            "y": "y:0",
+            "^y": "^y",
+            "z": "z:0",
+            "^z": "^z",
+            "foo_1:d:0": "foo_1:0",
+            "foo_1:e:0": "foo_1:1",
+            "^foo_1": "^foo_1",
+            "list_output:a:0": "list_output:0",
+            "list_output:a:1": "list_output:1",
+            "^list_output": "^list_output",
+            "foo_2:d:0": "foo_2:0",
+            "foo_2:e:0": "foo_2:1",
+            "^foo_2": "^foo_2",
+        })
+
+  def testShapes(self):
+    fdef = self._build_function_def()
+    g, _ = function_def_to_graph.function_def_to_graph_def(
+        fdef,
+        input_shapes=[
+            tensor_shape.TensorShape([]),
+            tensor_shape.TensorShape([5]), None
+        ])
+    self.assertEqual("shape" in g.node[0].attr, True)
+    self.assertSequenceEqual(
+        tensor_shape.TensorShape(g.node[0].attr["shape"].shape).as_list(), [])
+    self.assertEqual(g.node[0].attr["shape"].shape.unknown_rank, False)
+    self.assertEqual("shape" in g.node[1].attr, True)
+    self.assertSequenceEqual(
+        tensor_shape.TensorShape(g.node[1].attr["shape"].shape).as_list(), [5])
+    self.assertEqual(g.node[0].attr["shape"].shape.unknown_rank, False)
+    self.assertFalse("shape" in g.node[2].attr)
+
+  def testControlDependencies(self):
+
+    v = variables.Variable(1)
+
+    @def_function.function
+    def fn(inp):
+      assign = v.assign(3, name="assign", read_value=False)
+      x = constant_op.constant(2.0, name="x")
+      # TODO(b/79881896): Test external control dependency once that's
+      # supported.
+      with ops.control_dependencies([x, inp, assign]):
+        constant_op.constant(3.0, name="y")
+      return 4.0
+
+    inp = constant_op.constant(1.0)
+    fdef = fn.get_concrete_function(inp).function_def
+    func_graph = function_def_to_graph.function_def_to_graph(fdef)
+
+    op = func_graph.get_operation_by_name("y")
+    self.assertEqual(len(op.control_inputs), 3)
+    self.assertEqual(op.control_inputs[0].name, "assign")
+    self.assertEqual(op.control_inputs[1].name, "inp")
+    self.assertEqual(op.control_inputs[2].name, "x")
+
+  def testAttributesForArgDef(self):
+
+    @def_function.function
+    def fn(x):
+      return x
+
+    inp = constant_op.constant(1.0)
+    fdef = fn.get_concrete_function(inp).function_def
+    fdef.arg_attr[0].attr["_test_attr"].s = "value".encode("ascii")
+    graph_def = function_def_to_graph.function_def_to_graph_def(fdef)
+    placeholders = [
+        ndef for ndef in graph_def[0].node if ndef.op == "Placeholder"
+    ]
+    self.assertEqual(1, len(placeholders))
+    self.assertEqual(placeholders[0].attr["_test_attr"].s,
+                     "value".encode("ascii"))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/framework/gpu_util.py

@@ -0,0 +1,53 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Contains GPU utility functions."""
+
+import collections
+import re
+
+
+# Matches the DeviceAttributes.physical_device_desc field.
+_PHYSICAL_DEVICE_DESCRIPTION_REGEX = re.compile(
+    r'name: ([^,]*), (?:.*compute capability: (\d+)\.(\d+))?')
+
+
+# compute_capability is a (major version, minor version) pair, or None if this
+# is not an Nvidia GPU.
+GpuInfo = collections.namedtuple('gpu_info', ['name', 'compute_capability'])
+
+
+def compute_capability_from_device_desc(device_attrs):
+  """Returns the GpuInfo given a DeviceAttributes proto.
+
+  Args:
+    device_attrs: A DeviceAttributes proto.
+
+  Returns
+    A gpu_info tuple. Both fields are None if `device_attrs` does not have a
+    valid physical_device_desc field.
+  """
+  # TODO(jingyue): The device description generator has to be in sync with
+  # this file. Another option is to put compute capability in
+  # DeviceAttributes, but I avoided that to keep DeviceAttributes
+  # target-independent. Reconsider this option when we have more things like
+  # this to keep in sync.
+  # LINT.IfChange
+  match = _PHYSICAL_DEVICE_DESCRIPTION_REGEX.search(
+      device_attrs.physical_device_desc)
+  # LINT.ThenChange(//tensorflow/core/common_runtime/gpu/gpu_device.cc)
+  if not match:
+    return GpuInfo(None, None)
+  cc = (int(match.group(2)), int(match.group(3))) if match.group(2) else None
+  return GpuInfo(match.group(1), cc)

tensorflow/python/framework/graph_building_benchmark.py

@@ -0,0 +1,97 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+r"""Benchmarks for low-level graph building primitives.
+
+To run CPU benchmarks:
+  bazel run -c opt graph_building_benchmarks -- --benchmark_filter=.
+
+To run GPU benchmarks:
+  bazel run --config=cuda -c opt --copt="-mavx" graph_building_benchmarks -- \
+    --benchmark_filter=.
+
+To run a subset of benchmarks using --benchmarks flag.
+--benchmarks: the list of benchmarks to run. The specified value is interpreted
+as a regular expression and any benchmark whose name contains a partial match
+to the regular expression is executed.
+e.g. --benchmark_filter=".*MatMul.*" will run all matmul related benchmarks.
+
+"""
+import time
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.platform import test
+
+
+def run_benchmark(func, num_iters):
+  start = time.time()
+  for _ in range(num_iters):
+    func()
+  end = time.time()
+  return end - start
+
+
+class SingleOpBenchmarks(test.Benchmark):
+  """Benchmark for graph building time of ops."""
+
+  def _run_and_report(self, func, num_iters):
+    total_time = run_benchmark(func, num_iters)
+    mean_us = total_time * 1e6 / num_iters
+    self.report_benchmark(
+        iters=num_iters,
+        wall_time=mean_us,
+        extras={
+            "examples_per_sec": float("{0:.3f}".format(num_iters / total_time)),
+        })
+
+  def benchmarkAddScalars(self):
+    with context.execution_mode(context.GRAPH_MODE):
+      x = array_ops.placeholder(shape=[], dtype=dtypes.float32, name="x")
+      y = array_ops.placeholder(shape=[], dtype=dtypes.float32, name="y")
+
+      def bench():
+        return gen_math_ops.add(x, y)
+
+      self._run_and_report(bench, 1000)
+
+  def benchmarkAddBatchedMatrices(self):
+    with context.execution_mode(context.GRAPH_MODE):
+      x = array_ops.placeholder(
+          shape=[32, 784, 1000], dtype=dtypes.float32, name="x")
+      y = array_ops.placeholder(
+          shape=[32, 784, 1000], dtype=dtypes.float32, name="y")
+
+      def bench():
+        return gen_math_ops.add(x, y)
+
+      self._run_and_report(bench, 1000)
+
+  def benchmarkMatMul(self):
+    with context.execution_mode(context.GRAPH_MODE):
+      x = array_ops.placeholder(
+          shape=[784, 1000], dtype=dtypes.float32, name="x")
+      y = array_ops.placeholder(
+          shape=[1000, 1000], dtype=dtypes.float32, name="y")
+
+      def bench():
+        return gen_math_ops.mat_mul(x, y)
+
+      self._run_and_report(bench, 1000)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/framework/graph_io.py

@@ -0,0 +1,84 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utility functions for reading/writing graphs."""
+import os
+import os.path
+import sys
+
+from google.protobuf import text_format
+from tensorflow.python.framework import byte_swap_tensor
+from tensorflow.python.framework import ops
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('io.write_graph', v1=['io.write_graph', 'train.write_graph'])
+def write_graph(graph_or_graph_def, logdir, name, as_text=True):
+  """Writes a graph proto to a file.
+
+  The graph is written as a text proto unless `as_text` is `False`.
+
+  ```python
+  v = tf.Variable(0, name='my_variable')
+  sess = tf.compat.v1.Session()
+  tf.io.write_graph(sess.graph_def, '/tmp/my-model', 'train.pbtxt')
+  ```
+
+  or
+
+  ```python
+  v = tf.Variable(0, name='my_variable')
+  sess = tf.compat.v1.Session()
+  tf.io.write_graph(sess.graph, '/tmp/my-model', 'train.pbtxt')
+  ```
+
+  Args:
+    graph_or_graph_def: A `Graph` or a `GraphDef` protocol buffer.
+    logdir: Directory where to write the graph. This can refer to remote
+      filesystems, such as Google Cloud Storage (GCS).
+    name: Filename for the graph.
+    as_text: If `True`, writes the graph as an ASCII proto.
+
+  Returns:
+    The path of the output proto file.
+  """
+  if isinstance(graph_or_graph_def, ops.Graph):
+    graph_def = graph_or_graph_def.as_graph_def()
+  else:
+    graph_def = graph_or_graph_def
+
+  if sys.byteorder == 'big':
+    if hasattr(graph_def, 'node'):
+      byte_swap_tensor.swap_tensor_content_in_graph_node(
+          graph_def, 'big', 'little'
+      )
+    else:
+      byte_swap_tensor.swap_tensor_content_in_graph_function(
+          graph_def, 'big', 'little'
+      )
+
+  # gcs does not have the concept of directory at the moment.
+  if not logdir.startswith('gs:'):
+    file_io.recursive_create_dir(logdir)
+  path = os.path.join(logdir, name)
+  if as_text:
+    file_io.atomic_write_string_to_file(
+        path, text_format.MessageToString(graph_def)
+    )
+  else:
+    file_io.atomic_write_string_to_file(
+        path, graph_def.SerializeToString(deterministic=True))
+  return path

tensorflow/python/framework/graph_to_function_def.py

@@ -0,0 +1,183 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Utility to convert a Graph to a FunctionDef."""
+
+import re
+
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.framework import op_def_pb2
+from tensorflow.python.framework import op_def_registry
+
+
+def _make_argname_from_tensor_name(name):
+  return re.sub(":0$", "", name).replace(":", "_o")
+
+
+def _tensor_to_argdef(t, name=None, used_names=None):
+  """Convert tensor t to an argdef, with a specified name or a unique name."""
+  arg = op_def_pb2.OpDef.ArgDef()
+  if name is None:
+    arg.name = _make_argname_from_tensor_name(t.name)
+    if used_names is not None:
+      if arg.name in used_names:
+        i = 0
+        while True:
+          new_name = "%s_U%d" % (arg.name, i)
+          if new_name not in used_names:
+            arg.name = new_name
+            break
+          i += 1
+      used_names.add(arg.name)
+  else:
+    arg.name = name
+  arg.type = t.dtype.as_datatype_enum
+  return arg
+
+
+def _is_in_placeholders(op, func_arg_placeholders):
+  """Checks whether any output of this op is in func_arg_placeholders."""
+  return op.values() and any(x.name in func_arg_placeholders
+                             for x in op.values())
+
+
+def _get_node_def(op):
+  return op.node_def  # pylint: disable=protected-access
+
+
+def _get_op_def(op):
+  return op.op_def or op_def_registry.get(op.type)
+
+
+def _create_input_dict(function_graph,
+                       func_arg_placeholders,
+                       initial_value=None):
+  """Create a mapping from graph tensor names to function tensor names."""
+  if initial_value is None:
+    input_dict = {}
+  else:
+    input_dict = dict(initial_value)
+  for op in function_graph.get_operations():
+    if _is_in_placeholders(op, func_arg_placeholders):
+      input_dict[op.name] = op.name
+    else:
+      op_def = _get_op_def(op)
+      attrs = _get_node_def(op).attr
+      o = 0
+      for arg_def in op_def.output_arg:
+        if arg_def.number_attr:
+          num = attrs[arg_def.number_attr].i
+        elif arg_def.type_list_attr:
+          num = len(attrs[arg_def.type_list_attr].list.type)
+        else:
+          num = 1
+        for i in range(num):
+          result = "%s:%s:%d" % (op.name, arg_def.name, i)
+          input_dict[op.values()[o].name] = result
+          if o == 0:
+            input_dict[op.name] = result
+          o += 1
+  return input_dict
+
+
+def _add_op_node(op, func, input_dict):
+  """Converts an op to a function def node and add it to `func`."""
+  # Add an entry in func.node_def
+
+  # Note that extend() makes a copy in this case, see:
+  # https://developers.google.com/protocol-buffers/docs/reference/python-generated#repeated-message-fields
+  func.node_def.extend([_get_node_def(op)])
+  node_def = func.node_def[-1]
+  for i in range(len(node_def.input)):
+    if not node_def.input[i].startswith("^"):
+      assert node_def.input[i] in input_dict, ("%s missing from %s" %
+                                               (node_def.input[i],
+                                                input_dict.items()))
+      node_def.input[i] = input_dict[node_def.input[i]]
+  # The function is stateful if any of its operations are stateful.
+  # NOTE(mrry): The "Const" node typically does not have an `OpDef` associated
+  # with it, so we assume any nodes without an `OpDef` are stateless.
+  # TODO(skyewm): Remove the `is not None` test after we transition to the C
+  # API.
+  if op.op_def is not None and op.op_def.is_stateful:
+    func.signature.is_stateful = True
+
+
+def graph_to_function_def(graph, operations, inputs, outputs, out_names=None):
+  """Returns `graph` as a `FunctionDef` protocol buffer.
+
+  This method creates a [`FunctionDef`](
+  https://www.tensorflow.org/code/tensorflow/core/framework/function.proto)
+  protocol buffer that contains all the ops in `operations`.  The
+  operations become the body of the function.
+
+  The arguments `inputs` and `outputs` will be listed as the inputs
+  and outputs tensors of the function.  They must be lists of
+  tensors present in the graph.  The lists can optionally be empty.
+
+  Args:
+    graph: Graph.
+    operations: the operations to put in the function. Must be a subset of
+     the operations in the graph.
+    inputs: List of tensors. Inputs to the function.
+    outputs: List of tensors. Outputs of the function.
+    out_names: Optional list of string names for the outputs.
+
+  Returns:
+    A FunctionDef protocol buffer.
+
+  Raises:
+    ValueError: if out_names is specified and the wrong length.
+  """
+  func = function_pb2.FunctionDef()
+  func.signature.name = "_"
+  used_names = set()
+  func.signature.input_arg.extend(
+      [_tensor_to_argdef(i, used_names=used_names) for i in inputs])
+  # Initializes the input map with all placeholder input tensors.
+  initial_dict = {}
+  for o, m in zip(inputs, func.signature.input_arg):
+    initial_dict[o.name] = m.name
+  if out_names is None:
+    used_names = set()
+    func.signature.output_arg.extend(
+        [_tensor_to_argdef(o, used_names=used_names) for o in outputs])
+  elif len(outputs) != len(out_names):
+    raise ValueError(
+        f"out_names must be either empty or equal in size to outputs. "
+        f"len(out_names) = {len(out_names)} len(outputs) = {len(outputs)}")
+  elif len(out_names) != len(set(out_names)):
+    raise ValueError(
+        f"Must not have duplicates in out_names. Received: {out_names}")
+  else:
+    func.signature.output_arg.extend(
+        [_tensor_to_argdef(o, name=n) for o, n in zip(outputs, out_names)])
+  func_arg_placeholders = set(i.name for i in inputs)
+  input_dict = _create_input_dict(graph, func_arg_placeholders,
+                                  initial_value=initial_dict)
+
+  for op in operations:
+    if _is_in_placeholders(op, func_arg_placeholders):
+      continue
+    _add_op_node(op, func, input_dict)
+
+  if out_names is None:
+    for index, o in enumerate(outputs):
+      k = func.signature.output_arg[index].name
+      func.ret[k] = input_dict[o.name]
+  else:
+    for o, n in zip(outputs, out_names):
+      func.ret[n] = input_dict[o.name]
+
+  return func

tensorflow/python/framework/graph_util.py

@@ -0,0 +1,28 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Helpers to manipulate a tensor graph in python.
+
+API docstring: tensorflow.graph_util
+"""
+
+
+# pylint: disable=unused-import
+from tensorflow.python.framework.graph_util_impl import extract_sub_graph
+from tensorflow.python.framework.graph_util_impl import graph_defs_equal
+from tensorflow.python.framework.graph_util_impl import must_run_on_cpu
+from tensorflow.python.framework.graph_util_impl import remove_training_nodes
+from tensorflow.python.framework.graph_util_impl import tensor_shape_from_node_def_name
+# pylint: enable=unused-import

tensorflow/python/framework/graph_util_impl.py

@@ -0,0 +1,432 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Helpers to manipulate a tensor graph in python.
+"""
+
+import copy
+import re
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.framework import node_def_pb2
+from tensorflow.python.framework import _proto_comparators
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+GraphDef = tf_export(v1=["GraphDef"])(graph_pb2.GraphDef)
+
+_VARIABLE_OPS = {
+    "Assign",
+    "AssignAdd",
+    "AssignSub",
+    "Queue",
+    "ScatterAdd",
+    "ScatterSub",
+    "ScatterUpdate",
+    "TruncatedNormal",
+    "Variable",
+    "VariableV2",
+}
+
+_CONTROL_FLOW_OP_NAMES_OR_IDENTITY = [
+    "Switch",
+    "Enter",
+    "Exit",
+    "Identity",
+    "Merge",
+    "NextIteration",
+]
+
+_DEPRECATION_MSG = (
+    "This API was designed for TensorFlow v1. See "
+    "https://www.tensorflow.org/guide/migrate for instructions on how to "
+    "migrate your code to TensorFlow v2.")
+
+
+def _is_variable_op(op):
+  """Returns true if 'op' refers to a Variable node."""
+  return op in _VARIABLE_OPS
+
+# GraphDef protobuf docstring.
+graph_pb2.GraphDef.__doc__ = """\
+A protobuf containing the graph of operations.
+
+@compatibility(TF2)
+This API is not available in TensorFlow 2.x.
+
+You should not need to use `GraphDef`s directly in TF2. To load `GraphDef`s in
+TF2, use SavedModel. The SavedModel contains the `GraphDef`.
+
+Before:
+
+```python
+with tf.io.gfile.GFile('/tmp/graph.pb', 'rb') as f:
+  graph_def = tf.compat.v1.GraphDef()
+  graph_def.ParseFromString(f.read())
+```
+
+After:
+
+```python
+tf.saved_model.load('/tmp/saved_model')
+```
+
+If you would like to create a `GraphDef` in TF2, use `tf.function` and
+`get_concrete_function`.
+
+>>> @tf.function
+>>> def f(x):
+>>>   return x
+>>>
+>>> graph_def = f.get_concrete_function(1.).graph.as_graph_def()
+>>> print(graph_def)
+
+@end_compatibility
+
+"""
+
+
+@deprecation.deprecated(
+    date=None,
+    instructions=_DEPRECATION_MSG)
+@tf_export(v1=["graph_util.must_run_on_cpu"])
+def must_run_on_cpu(node, pin_variables_on_cpu=False):
+  """Returns True if the given node_def must run on CPU, otherwise False.
+
+  Args:
+    node: The node to be assigned to a device. Could be either an ops.Operation
+      or NodeDef.
+    pin_variables_on_cpu: If True, this function will return False if node_def
+      represents a variable-related op.
+
+  Returns:
+    True if the given node must run on CPU, otherwise False.
+  """
+
+  if isinstance(node, ops.Operation):
+    node_def = node.node_def
+  else:
+    assert isinstance(node, node_def_pb2.NodeDef)
+    node_def = node
+
+  # If the op is a variable-related op, should we pin it on CPU?
+  if pin_variables_on_cpu and _is_variable_op(node_def.op):
+    return True
+
+  # Constant operations producing a string or int32 must run on CPU.
+  if node_def.op == "Const":
+    # Get the value of the 'dtype' attr
+    dtype = node_def.attr["dtype"].type
+    if dtype == dtypes.string or dtype == dtypes.int32:
+      return True
+
+  if node_def.op in ["DynamicStitch", "ParallelDynamicStitch"]:
+    dtype = node_def.attr["T"].type
+    if dtype == dtypes.int32:
+      # DynamicStitch on GPU only works for int32 values.
+      return True
+
+  if node_def.op in ["Cast"]:
+    dtype = node_def.attr["SrcT"].type
+    if dtype == dtypes.int32:
+      # Cast on GPU does not works for int32 values.
+      return True
+  return False
+
+
+################################################################################
+#
+# device functions for use in with g.device(...)
+#
+################################################################################
+
+
+def _node_name(n):
+  if n.startswith("^"):
+    return n[1:]
+  else:
+    return n.split(":")[0]
+
+
+def _get_colocated_node_name(colocated_node_name):
+  """Decodes colocated node name and returns it without loc:@ prepended."""
+  colocated_node_decoded = colocated_node_name.decode("utf-8")
+  if colocated_node_decoded.startswith("loc:@"):
+    return colocated_node_decoded[5:]
+  return colocated_node_decoded
+
+
+def _extract_graph_summary(graph_def):
+  """Extracts useful information from the graph and returns them."""
+  name_to_input_name = {}  # Keyed by the dest node name.
+  name_to_node = {}  # Keyed by node name.
+
+  # Keeps track of node sequences. It is important to still output the
+  # operations in the original order.
+  name_to_seq_num = {}  # Keyed by node name.
+  seq = 0
+  for node in graph_def.node:
+    n = _node_name(node.name)
+    name_to_node[n] = node
+    name_to_input_name[n] = [_node_name(x) for x in node.input]
+    # Prevent colocated nodes from being lost.
+    if "_class" in node.attr:
+      for colocated_node_name in node.attr["_class"].list.s:
+        name_to_input_name[n].append(
+            _get_colocated_node_name(colocated_node_name))
+    name_to_seq_num[n] = seq
+    seq += 1
+  return name_to_input_name, name_to_node, name_to_seq_num
+
+
+def _assert_nodes_are_present(name_to_node, nodes):
+  """Assert that nodes are present in the graph."""
+  for d in nodes:
+    assert d in name_to_node, "%s is not in graph" % d
+
+
+def _bfs_for_reachable_nodes(target_nodes, name_to_input_name):
+  """Breadth first search for reachable nodes from target nodes."""
+  nodes_to_keep = set()
+  # Breadth first search to find all the nodes that we should keep.
+  next_to_visit = list(target_nodes)
+  while next_to_visit:
+    node = next_to_visit[0]
+    del next_to_visit[0]
+    if node in nodes_to_keep:
+      # Already visited this node.
+      continue
+    nodes_to_keep.add(node)
+    if node in name_to_input_name:
+      next_to_visit += name_to_input_name[node]
+  return nodes_to_keep
+
+
+@deprecation.deprecated(
+    date=None,
+    instructions=_DEPRECATION_MSG)
+@tf_export(v1=["graph_util.extract_sub_graph"])
+def extract_sub_graph(graph_def, dest_nodes):
+  """Extract the subgraph that can reach any of the nodes in 'dest_nodes'.
+
+  Args:
+    graph_def: A graph_pb2.GraphDef proto.
+    dest_nodes: An iterable of strings specifying the destination node names.
+  Returns:
+    The GraphDef of the sub-graph.
+
+  Raises:
+    TypeError: If 'graph_def' is not a graph_pb2.GraphDef proto.
+  """
+
+  if not isinstance(graph_def, graph_pb2.GraphDef):
+    raise TypeError("graph_def must be a graph_pb2.GraphDef proto, but got "
+                    f"type {type(graph_def)}.")
+
+  if isinstance(dest_nodes, str):
+    raise TypeError("dest_nodes must be an iterable of strings, but got "
+                    f"type {type(dest_nodes)}.")
+
+  name_to_input_name, name_to_node, name_to_seq_num = _extract_graph_summary(
+      graph_def)
+  _assert_nodes_are_present(name_to_node, dest_nodes)
+
+  nodes_to_keep = _bfs_for_reachable_nodes(dest_nodes, name_to_input_name)
+
+  nodes_to_keep_list = sorted(
+      list(nodes_to_keep), key=lambda n: name_to_seq_num[n])
+  # Now construct the output GraphDef
+  out = graph_pb2.GraphDef()
+  for n in nodes_to_keep_list:
+    out.node.extend([copy.deepcopy(name_to_node[n])])
+  out.library.CopyFrom(graph_def.library)
+  out.versions.CopyFrom(graph_def.versions)
+
+  return out
+
+
+@deprecation.deprecated(
+    date=None,
+    instructions=_DEPRECATION_MSG)
+@tf_export(v1=["graph_util.tensor_shape_from_node_def_name"])
+def tensor_shape_from_node_def_name(graph, input_name):
+  """Convenience function to get a shape from a NodeDef's input string."""
+  # To get a tensor, the name must be in the form <input>:<port>, for example
+  # 'Mul:0'. The GraphDef input strings don't always have the port specified
+  # though, so if there isn't a colon we need to add a default ':0' to the end.
+  if ":" not in input_name:
+    canonical_name = input_name + ":0"
+  else:
+    canonical_name = input_name
+  tensor = graph.get_tensor_by_name(canonical_name)
+  shape = tensor.get_shape()
+  return shape
+
+
+@deprecation.deprecated(
+    date=None,
+    instructions=_DEPRECATION_MSG)
+@tf_export(v1=["graph_util.remove_training_nodes"])
+def remove_training_nodes(input_graph, protected_nodes=None):
+  """Prunes out nodes that aren't needed for inference.
+
+  There are nodes like Identity and CheckNumerics that are only useful
+  during training, and can be removed in graphs that will be used for
+  nothing but inference. Here we identify and remove them, returning an
+  equivalent graph. To be specific, CheckNumerics nodes are always removed, and
+  Identity nodes that aren't involved in control edges are spliced out so that
+  their input and outputs are directly connected.
+
+  Args:
+    input_graph: Model to analyze and prune.
+    protected_nodes: An optional list of names of nodes to be kept
+      unconditionally. This is for example useful to preserve Identity output
+      nodes.
+
+  Returns:
+    A list of nodes with the unnecessary ones removed.
+  """
+  if not protected_nodes:
+    protected_nodes = []
+
+  types_to_remove = {"CheckNumerics": True}
+
+  input_nodes = input_graph.node
+  names_to_remove = {}
+  for node in input_nodes:
+    if node.op in types_to_remove and node.name not in protected_nodes:
+      names_to_remove[node.name] = True
+
+  nodes_after_removal = []
+  for node in input_nodes:
+    if node.name in names_to_remove:
+      continue
+    new_node = node_def_pb2.NodeDef()
+    new_node.CopyFrom(node)
+    input_before_removal = node.input
+    del new_node.input[:]
+    for full_input_name in input_before_removal:
+      input_name = re.sub(r"^\^", "", full_input_name)
+      if input_name in names_to_remove:
+        continue
+      new_node.input.append(full_input_name)
+    nodes_after_removal.append(new_node)
+
+  types_to_splice = {"Identity": True, "StopGradient": True}
+  control_input_names = set()
+  node_names_with_control_input = set()
+  node_in_colocated = set()
+
+  for node in nodes_after_removal:
+    for node_input in node.input:
+      if "^" in node_input:
+        control_input_names.add(node_input.replace("^", ""))
+        node_names_with_control_input.add(node.name)
+    # Prevent colocated nodes from being lost.
+    if "_class" in node.attr:
+      for colocated_node_name in node.attr["_class"].list.s:
+        node_in_colocated.add(_get_colocated_node_name(colocated_node_name))
+
+  names_to_splice = {}
+  for node in nodes_after_removal:
+    if node.op in types_to_splice and node.name not in protected_nodes:
+      if node.name in node_in_colocated:
+        continue
+      # We don't want to remove nodes that have control edge inputs, because
+      # they might be involved in subtle dependency issues that removing them
+      # will jeopardize.
+      if node.name not in node_names_with_control_input:
+        names_to_splice[node.name] = node.input[0]
+
+  # We also don't want to remove nodes which are used as control edge inputs.
+  names_to_splice = {name: value for name, value in names_to_splice.items()
+                     if name not in control_input_names}
+
+  nodes_after_splicing = []
+  for node in nodes_after_removal:
+    if node.name in names_to_splice:
+      continue
+    new_node = node_def_pb2.NodeDef()
+    new_node.CopyFrom(node)
+    input_before_removal = node.input
+    del new_node.input[:]
+    for full_input_name in input_before_removal:
+      input_name = re.sub(r"^\^", "", full_input_name)
+      while input_name in names_to_splice:
+        full_input_name = names_to_splice[input_name]
+        input_name = re.sub(r"^\^", "", full_input_name)
+      new_node.input.append(full_input_name)
+    nodes_after_splicing.append(new_node)
+
+  output_graph = graph_pb2.GraphDef()
+  output_graph.node.extend(nodes_after_splicing)
+  return output_graph
+
+
+@tf_export("__internal__.graph_util.graph_defs_equal", v1=[])
+def graph_defs_equal(graph_def_1: graph_pb2.GraphDef,
+                     graph_def_2: graph_pb2.GraphDef,
+                     treat_nan_as_equal: bool = False) -> bool:
+  """Returns True iff the graph def arguments are structurally equivalent.
+
+  The notion of equivalence encoded here checks that the set of NodeDefs in
+  the GraphDef's function library and main graph body are identical.
+  Additionally, it checks that the functions in the function library are equal
+  as sets.
+
+  Example usage:
+
+  ```
+  with tf.Graph().as_default() as g1:
+    tf.constant(1)
+
+  with tf.Graph().as_default() as g2:
+    tf.constant(2)
+
+  with tf.Graph().as_default() as g3:
+    tf.constant(1)
+
+  assert tf.__internal__.graph_util.graph_defs_equal(g1.as_graph_def(),
+                                                     g3.as_graph_def())
+
+  assert not tf.__internal__.graph_util.graph_defs_equal(g1.as_graph_def(),
+                                                         g2.as_graph_def())
+  ```
+
+  Args:
+    graph_def_1: Instance of `graph_pb2.GraphDef` to compare.
+    graph_def_2: Instance of `graph_pb2.GraphDef` to compare.
+    treat_nan_as_equal: Boolean indicating whether or not to treat nan
+      floating-point values as equal. This is crucial for any equivalence
+      relation defined over GraphDefs, to ensure symmetry.
+
+  Returns:
+    Boolean indicating structural equivalence as described above.
+
+  Raises:
+    TypeError: If either of the GraphDefs are not instances of
+      `graph_pb2.GraphDef`.
+  """
+  if not isinstance(graph_def_1, graph_pb2.GraphDef):
+    raise TypeError("graph_def_1 must be a graph_pb2.GraphDef proto, but got "
+                    f"type {type(graph_def_1)}.")
+  if not isinstance(graph_def_2, graph_pb2.GraphDef):
+    raise TypeError("graph_def_2 must be a graph_pb2.GraphDef proto, but got "
+                    f"type {type(graph_def_2)}.")
+  options = _proto_comparators.ProtoComparisonOptions(treat_nan_as_equal)
+  return _proto_comparators.EqualsGraphDef(graph_def_1.SerializeToString(),
+                                           graph_def_2.SerializeToString(),
+                                           options)

tensorflow/python/framework/graph_util_test.py

@@ -0,0 +1,502 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensorflow.python.client.graph_util."""
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.framework import node_def_pb2
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import graph_util
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gen_state_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+# Utility device function to use for testing
+def TestDeviceFuncPinVariableToCpu(op):
+  if op.device:
+    return op.device
+  return "/cpu:0" if op.node_def.op in ["Variable", "VariableV2"] else op.device
+
+
+class GraphUtilTest(test.TestCase):
+
+  def testTwoDeviceFunctions(self):
+    with ops.Graph().as_default() as g:
+      var_0 = gen_state_ops.variable(
+          shape=[1],
+          dtype=dtypes.float32,
+          name="var_0",
+          container="",
+          shared_name="")
+      with g.device(TestDeviceFuncPinVariableToCpu):
+        var_1 = gen_state_ops.variable(
+            shape=[1],
+            dtype=dtypes.float32,
+            name="var_1",
+            container="",
+            shared_name="")
+      var_2 = gen_state_ops.variable(
+          shape=[1],
+          dtype=dtypes.float32,
+          name="var_2",
+          container="",
+          shared_name="")
+      var_3 = gen_state_ops.variable(
+          shape=[1],
+          dtype=dtypes.float32,
+          name="var_3",
+          container="",
+          shared_name="")
+      with g.device(TestDeviceFuncPinVariableToCpu):
+        var_4 = gen_state_ops.variable(
+            shape=[1],
+            dtype=dtypes.float32,
+            name="var_4",
+            container="",
+            shared_name="")
+        with g.device("/device:GPU:0"):
+          var_5 = gen_state_ops.variable(
+              shape=[1],
+              dtype=dtypes.float32,
+              name="var_5",
+              container="",
+              shared_name="")
+        var_6 = gen_state_ops.variable(
+            shape=[1],
+            dtype=dtypes.float32,
+            name="var_6",
+            container="",
+            shared_name="")
+
+    self.assertDeviceEqual(var_0.device, None)
+    self.assertDeviceEqual(var_1.device, "/device:CPU:0")
+    self.assertDeviceEqual(var_2.device, None)
+    self.assertDeviceEqual(var_3.device, None)
+    self.assertDeviceEqual(var_4.device, "/device:CPU:0")
+    self.assertDeviceEqual(var_5.device, "/device:GPU:0")
+    self.assertDeviceEqual(var_6.device, "/device:CPU:0")
+
+  @test_util.run_v1_only("b/120545219")
+  def testNestedDeviceFunctions(self):
+    with ops.Graph().as_default():
+      var_0 = variable_v1.VariableV1(0)
+      with ops.device(TestDeviceFuncPinVariableToCpu):
+        var_1 = variable_v1.VariableV1(1)
+        with ops.device(lambda op: "/device:GPU:0"):
+          var_2 = variable_v1.VariableV1(2)
+        with ops.device("/device:GPU:0"):  # Implicit merging device function.
+          var_3 = variable_v1.VariableV1(3)
+
+    self.assertDeviceEqual(var_0.device, None)
+    self.assertDeviceEqual(var_1.device, "/device:CPU:0")
+    self.assertDeviceEqual(var_2.device, "/device:GPU:0")
+    self.assertDeviceEqual(var_3.device, "/device:GPU:0")
+
+  def testExplicitDevice(self):
+    with ops.Graph().as_default() as g:
+      const_0 = constant_op.constant(5.0)
+      with g.device("/device:GPU:0"):
+        const_1 = constant_op.constant(5.0)
+      with g.device("/device:GPU:1"):
+        const_2 = constant_op.constant(5.0)
+      with g.device("/device:CPU:0"):
+        const_3 = constant_op.constant(5.0)
+      with g.device("/device:CPU:1"):
+        const_4 = constant_op.constant(5.0)
+      with g.device("/job:ps"):
+        const_5 = constant_op.constant(5.0)
+
+    self.assertDeviceEqual(const_0.device, None)
+    self.assertDeviceEqual(const_1.device, "/device:GPU:0")
+    self.assertDeviceEqual(const_2.device, "/device:GPU:1")
+    self.assertDeviceEqual(const_3.device, "/device:CPU:0")
+    self.assertDeviceEqual(const_4.device, "/device:CPU:1")
+    self.assertDeviceEqual(const_5.device, "/job:ps")
+
+  def testDefaultDevice(self):
+    with ops.Graph().as_default() as g, g.device(
+        TestDeviceFuncPinVariableToCpu):
+      with g.device("/job:ps"):
+        const_0 = constant_op.constant(5.0)
+      with g.device("/device:GPU:0"):
+        const_1 = constant_op.constant(5.0)
+      with g.device("/device:GPU:1"):
+        const_2 = constant_op.constant(5.0)
+      with g.device("/device:CPU:0"):
+        const_3 = constant_op.constant(5.0)
+      with g.device("/device:CPU:1"):
+        const_4 = constant_op.constant(5.0)
+      with g.device("/replica:0"):
+        const_5 = constant_op.constant(5.0)
+
+    self.assertDeviceEqual(const_0.device, "/job:ps")
+    self.assertDeviceEqual(const_1.device, "/device:GPU:0")
+    self.assertDeviceEqual(const_2.device, "/device:GPU:1")
+    self.assertDeviceEqual(const_3.device, "/device:CPU:0")
+    self.assertDeviceEqual(const_4.device, "/device:CPU:1")
+    self.assertDeviceEqual(const_5.device, "/replica:0")
+
+  def testExtractSubGraph(self):
+    graph_def = graph_pb2.GraphDef()
+    n1 = graph_def.node.add()
+    n1.name = "n1"
+    n1.input.extend(["n5"])
+    n2 = graph_def.node.add()
+    n2.name = "n2"
+    # Take the first output of the n1 node as the input.
+    n2.input.extend(["n1:0"])
+    n3 = graph_def.node.add()
+    n3.name = "n3"
+    # Add a control input (which isn't really needed by the kernel, but
+    # rather to enforce execution order between nodes).
+    n3.input.extend(["^n2"])
+    n4 = graph_def.node.add()
+    n4.name = "n4"
+
+    # It is fine to have a loops in the graph as well.
+    n5 = graph_def.node.add()
+    n5.name = "n5"
+    n5.input.extend(["n1"])
+
+    sub_graph = graph_util.extract_sub_graph(graph_def, ["n3"])
+    self.assertEqual("n1", sub_graph.node[0].name)
+    self.assertEqual("n2", sub_graph.node[1].name)
+    self.assertEqual("n3", sub_graph.node[2].name)
+    self.assertEqual("n5", sub_graph.node[3].name)
+
+  def testExtractSubGraphWithInvalidDestNodes(self):
+    graph_def = graph_pb2.GraphDef()
+    n1 = graph_def.node.add()
+    n1.name = "n1"
+    with self.assertRaisesRegex(TypeError, "must be an iterable"):
+      graph_util.extract_sub_graph(graph_def, "n1")
+
+  def create_node_def(self, op, name, inputs):
+    new_node = node_def_pb2.NodeDef()
+    new_node.op = op
+    new_node.name = name
+    new_node.input.extend(inputs)
+    return new_node
+
+  def create_constant_node_def(self,
+                               name,
+                               value,
+                               dtype,
+                               shape=None,
+                               inputs=None):
+    node = self.create_node_def("Const", name, inputs or [])
+    self.set_attr_dtype(node, "dtype", dtype)
+    self.set_attr_tensor(node, "value", value, dtype, shape)
+    return node
+
+  def set_attr_dtype(self, node, key, value):
+    node.attr[key].CopyFrom(
+        attr_value_pb2.AttrValue(type=value.as_datatype_enum))
+
+  def set_attr_list(self, node, key, value_list):
+    node.attr[key].CopyFrom(
+        attr_value_pb2.AttrValue(
+            list=attr_value_pb2.AttrValue.ListValue(s=value_list)
+        )
+    )
+
+  def set_attr_tensor(self, node, key, value, dtype, shape=None):
+    node.attr[key].CopyFrom(
+        attr_value_pb2.AttrValue(
+            tensor=tensor_util.make_tensor_proto(
+                value, dtype=dtype, shape=shape)))
+
+  def testRemoveTrainingNodes(self):
+    a_constant_name = "a_constant"
+    b_constant_name = "b_constant"
+    c_constant_name = "c_constant"
+    a_check_name = "a_check"
+    b_check_name = "b_check"
+    a_identity_name = "a_identity"
+    b_identity_name = "b_identity"
+    c_identity_name = "c_identity"
+    add_name = "add"
+    sub_name = "sub"
+    graph_def = graph_pb2.GraphDef()
+    a_constant = self.create_constant_node_def(
+        a_constant_name, value=1, dtype=dtypes.float32, shape=[])
+    graph_def.node.extend([a_constant])
+    a_check_node = self.create_node_def("CheckNumerics", a_check_name,
+                                        [a_constant_name])
+    graph_def.node.extend([a_check_node])
+    a_identity_node = self.create_node_def(
+        "Identity", a_identity_name, [a_constant_name, "^" + a_check_name])
+    graph_def.node.extend([a_identity_node])
+    b_constant = self.create_constant_node_def(
+        b_constant_name, value=1, dtype=dtypes.float32, shape=[])
+    graph_def.node.extend([b_constant])
+    b_check_node = self.create_node_def("CheckNumerics", b_check_name,
+                                        [b_constant_name])
+    graph_def.node.extend([b_check_node])
+    b_identity_node = self.create_node_def(
+        "Identity", b_identity_name, [b_constant_name, "^" + b_check_name])
+    graph_def.node.extend([b_identity_node])
+    add_node = self.create_node_def("Add", add_name,
+                                    [a_identity_name, b_identity_name])
+    self.set_attr_dtype(add_node, "T", dtypes.float32)
+    graph_def.node.extend([add_node])
+    c_constant = self.create_constant_node_def(
+        c_constant_name, value=1, dtype=dtypes.float32, shape=[]
+    )
+    graph_def.node.extend([c_constant])
+    c_identity_node = self.create_node_def(
+        "Identity", c_identity_name, [c_constant_name]
+    )
+    graph_def.node.extend([c_identity_node])
+
+    sub_node = self.create_node_def(
+        "Sub", sub_name, [c_constant_name, c_identity_name]
+    )
+    self.set_attr_list(sub_node, "_class", [compat.as_bytes(c_identity_name)])
+    graph_def.node.extend([sub_node])
+
+    expected_output = graph_pb2.GraphDef()
+    a_constant = self.create_constant_node_def(
+        a_constant_name, value=1, dtype=dtypes.float32, shape=[])
+    expected_output.node.extend([a_constant])
+    b_constant = self.create_constant_node_def(
+        b_constant_name, value=1, dtype=dtypes.float32, shape=[])
+    expected_output.node.extend([b_constant])
+    add_node = self.create_node_def("Add", add_name,
+                                    [a_constant_name, b_constant_name])
+    self.set_attr_dtype(add_node, "T", dtypes.float32)
+    expected_output.node.extend([add_node])
+    c_constant = self.create_constant_node_def(
+        c_constant_name, value=1, dtype=dtypes.float32, shape=[]
+    )
+    expected_output.node.extend([c_constant])
+    c_identity_node = self.create_node_def(
+        "Identity", c_identity_name, [c_constant_name]
+    )
+    expected_output.node.extend([c_identity_node])
+
+    sub_node = self.create_node_def(
+        "Sub", sub_name, [c_constant_name, c_identity_name]
+    )
+    self.set_attr_list(sub_node, "_class", [compat.as_bytes(c_identity_name)])
+    expected_output.node.extend([sub_node])
+
+    output = graph_util.remove_training_nodes(graph_def)
+    self.assertProtoEquals(expected_output, output)
+
+  def testRemoveIdentityChains(self):
+    """Check that chains of Identity nodes are correctly pruned.
+
+    Create a chain of four nodes, A, B, C, and D where A inputs B, B inputs C,
+    and C inputs D. Nodes B and C are "Identity" and should be pruned, resulting
+    in the nodes A and D, where A inputs D.
+    """
+    graph_def = graph_pb2.GraphDef()
+    graph_def.node.extend([
+        self.create_node_def("Aop", "A", ["B"]),
+        self.create_node_def("Identity", "B", ["C"]),
+        self.create_node_def("Identity", "C", ["D"]),
+        self.create_node_def("Dop", "D", [])
+    ])
+
+    expected_graph_def = graph_pb2.GraphDef()
+    expected_graph_def.node.extend([
+        self.create_node_def("Aop", "A", ["D"]),
+        self.create_node_def("Dop", "D", [])
+    ])
+
+    self.assertProtoEquals(expected_graph_def,
+                           graph_util.remove_training_nodes(graph_def))
+
+  def testRemoveIdentityUsedAsControlInputInConst(self):
+    """Check that Identity nodes used as control inputs are not removed."""
+    graph_def = graph_pb2.GraphDef()
+    graph_def.node.extend([
+        self.create_constant_node_def("C", 1, dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    self.assertProtoEquals(graph_def,
+                           graph_util.remove_training_nodes(graph_def))
+
+  def testSimpleGraphdefsCompareEqual(self):
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.node.extend([
+        self.create_constant_node_def("C", 1, dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.node.extend([
+        self.create_constant_node_def("C", 1, dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    self.assertTrue(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+  def testNodeDefsInDifferentOrderCompareEqual(self):
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.node.extend([
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", []),
+        self.create_constant_node_def("C", 1, dtypes.float32, inputs=["^I"]),
+    ])
+
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.node.extend([
+        self.create_constant_node_def("C", 1, dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    self.assertTrue(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+  def testDifferentGraphDefsCompareNotEqual(self):
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.node.extend([
+        self.create_constant_node_def("C", 1, dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.node.extend([
+        self.create_constant_node_def("C", 2, dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+    self.assertFalse(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+  def testGraphdefsWithNanCompareNonEqual(self):
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.node.extend([
+        self.create_constant_node_def(
+            "C", float("nan"), dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.node.extend([
+        self.create_constant_node_def(
+            "C", float("nan"), dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+    self.assertFalse(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+  def testSimpleGraphdefEqualityWithNansEqual(self):
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.node.extend([
+        self.create_constant_node_def(
+            "C", float("nan"), dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.node.extend([
+        self.create_constant_node_def(
+            "C", float("nan"), dtypes.float32, inputs=["^I"]),
+        self.create_node_def("Identity", "I", ["Base"]),
+        self.create_node_def("BaseOp", "Base", [])
+    ])
+    self.assertTrue(
+        graph_util.graph_defs_equal(
+            graph_def1, graph_def2, treat_nan_as_equal=True))
+
+  def testGraphDefsWithFunctionLibsCompareEqual(self):
+
+    @function.Defun(dtypes.float32)
+    def F1(x):
+      return math_ops.exp(x) - math_ops.exp(-x)
+
+    library = function_pb2.FunctionDefLibrary()
+    library.function.extend([F1.definition])
+
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.library.CopyFrom(library)
+
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.library.CopyFrom(library)
+
+    self.assertTrue(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+  def testGraphDefsWithPermutedFunctionsCompareEqual(self):
+
+    @function.Defun(dtypes.float32)
+    def F1(x):
+      return math_ops.exp(x) - math_ops.exp(-x)
+
+    @function.Defun(dtypes.float32)
+    def F2(x):
+      return math_ops.exp(x)
+
+    definition_1 = F1.definition
+    definition_2 = F2.definition
+    library = function_pb2.FunctionDefLibrary()
+    library.function.extend([definition_1, definition_2])
+
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.library.CopyFrom(library)
+
+    reversed_library = function_pb2.FunctionDefLibrary()
+    reversed_library.function.extend([definition_2, definition_1])
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.library.CopyFrom(reversed_library)
+
+    self.assertTrue(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+  def testGraphDefsWithPermutedNodesInFunctionsCompareEqual(self):
+
+    @function.Defun(dtypes.float32)
+    def F1(x):
+      return math_ops.exp(x) - math_ops.exp(-x)
+
+    f1_def = F1.definition
+
+    library = function_pb2.FunctionDefLibrary()
+    library.function.extend([f1_def])
+
+    graph_def1 = graph_pb2.GraphDef()
+    graph_def1.library.CopyFrom(library)
+
+    reversed_function = function_pb2.FunctionDef()
+    reversed_function.CopyFrom(f1_def)
+    # Clear the node_def attribute.
+    del reversed_function.node_def[:]
+    reversed_function.node_def.extend(reversed(f1_def.node_def))
+    reversed_library = function_pb2.FunctionDefLibrary()
+    reversed_library.function.extend([reversed_function])
+    graph_def2 = graph_pb2.GraphDef()
+    graph_def2.library.CopyFrom(reversed_library)
+
+    self.assertTrue(graph_util.graph_defs_equal(graph_def1, graph_def2))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/framework/immutable_dict.py

@@ -0,0 +1,47 @@
+# 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.
+# ==============================================================================
+"""Immutable mapping."""
+
+import collections.abc
+
+
+# WARNING: this class is used internally by extension types (tf.ExtensionType),
+# and may be deleted if/when extension types transition to a different encoding
+# in the future.
+class ImmutableDict(collections.abc.Mapping):
+  """Immutable `Mapping`."""
+  # Note: keys, items, values, get, __eq__, and __ne__ are implemented by
+  # the `Mapping` base class.
+
+  def __init__(self, *args, **kwargs):
+    self._dict = dict(*args, **kwargs)
+
+  def __getitem__(self, key):
+    return self._dict[key]
+
+  def __contains__(self, key):
+    return key in self._dict
+
+  def __iter__(self):
+    return iter(self._dict)
+
+  def __len__(self):
+    return len(self._dict)
+
+  def __repr__(self):
+    return f'ImmutableDict({self._dict})'
+
+  # This suppresses a warning that tf.nest would otherwise generate.
+  __supported_by_tf_nest__ = True

tensorflow/python/framework/immutable_dict_test.py

@@ -0,0 +1,98 @@
+# 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 tf.framework.immutable_dict."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import immutable_dict
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+
+class ImmutableDictTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  def testGetItem(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(d['x'], 1)
+    self.assertEqual(d['y'], 2)
+    with self.assertRaises(KeyError):
+      d['z']  # pylint: disable=pointless-statement
+
+  def testIter(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(set(iter(d)), set(['x', 'y']))
+
+  def testContains(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertIn('x', d)
+    self.assertIn('y', d)
+    self.assertNotIn('z', d)
+
+  def testLen(self):
+    d1 = immutable_dict.ImmutableDict({})
+    self.assertLen(d1, 0)  # pylint: disable=g-generic-assert
+    d2 = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertLen(d2, 2)
+
+  def testRepr(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    s = repr(d)
+    self.assertTrue(s == "ImmutableDict({'x': 1, 'y': 2})" or
+                    s == "ImmutableDict({'y': 1, 'x': 2})")
+
+  def testGet(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(d.get('x'), 1)
+    self.assertEqual(d.get('y'), 2)
+    self.assertIsNone(d.get('z'))
+    self.assertEqual(d.get('z', 'Foo'), 'Foo')
+
+  def testKeys(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(set(d.keys()), set(['x', 'y']))
+
+  def testValues(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(set(d.values()), set([1, 2]))
+
+  def testItems(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(set(d.items()), set([('x', 1), ('y', 2)]))
+
+  def testEqual(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertEqual(d, {'x': 1, 'y': 2})
+
+  def testNotEqual(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    self.assertNotEqual(d, {'x': 1})
+
+  def testSetItemFails(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    with self.assertRaises(TypeError):
+      d['x'] = 5  # pylint: disable=unsupported-assignment-operation
+    with self.assertRaises(TypeError):
+      d['z'] = 5  # pylint: disable=unsupported-assignment-operation
+
+  def testDelItemFails(self):
+    d = immutable_dict.ImmutableDict({'x': 1, 'y': 2})
+    with self.assertRaises(TypeError):
+      del d['x']  # pylint: disable=unsupported-delete-operation
+    with self.assertRaises(TypeError):
+      del d['z']  # pylint: disable=unsupported-delete-operation
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/framework/importer.py

@@ -0,0 +1,566 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A utility function for importing TensorFlow graphs."""
+import contextlib
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.python import tf2
+from tensorflow.python.client import pywrap_tf_session as c_api
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import function
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.util import compat
+from tensorflow.python.util.deprecation import deprecated_args
+from tensorflow.python.util.tf_export import tf_export
+
+
+# TODO(b/307794935): Remove after bug is fixed.
+is_oss = True  # Updated by copybara.
+
+
+def _IsControlInput(input_name):
+  # Expected format: '^operation_name' (control input).
+  return input_name.startswith('^')
+
+
+def _ParseTensorName(tensor_name):
+  """Parses a tensor name into an operation name and output index.
+
+  This function will canonicalize tensor names as follows:
+
+  * "foo:0"       -> ("foo", 0)
+  * "foo:7"       -> ("foo", 7)
+  * "foo"         -> ("foo", 0)
+  * "foo:bar:baz" -> ValueError
+
+  Args:
+    tensor_name: The name of a tensor.
+
+  Returns:
+    A tuple containing the operation name, and the output index.
+
+  Raises:
+    ValueError: If `tensor_name' cannot be interpreted as the name of a tensor.
+  """
+  components = tensor_name.split(':')
+  if len(components) == 2:
+    # Expected format: 'operation_name:output_index'.
+    try:
+      output_index = int(components[1])
+    except ValueError:
+      raise ValueError(f'Cannot convert {tensor_name!r} to a tensor name. '
+                       'Second component of the name following the `:` should '
+                       f'be an int. Got {components[1]}.')
+    return components[0], output_index
+  elif len(components) == 1:
+    # Expected format: 'operation_name' (implicit 0th output).
+    return components[0], 0
+  else:
+    raise ValueError(f"Cannot convert '{tensor_name}' to a tensor name. Tensor "
+                     'names should not contain more than 1 `:`. Obtained '
+                     f'{len(components) - 1}')
+
+
+@contextlib.contextmanager
+def _MaybeDevice(device):
+  """Applies the given device only if device is not None or empty."""
+  if device:
+    with ops.device(device):
+      yield
+  else:
+    yield
+
+
+def _ProcessGraphDefParam(graph_def):
+  """Type-checks and possibly canonicalizes `graph_def`."""
+  if not isinstance(graph_def, graph_pb2.GraphDef):
+    # `graph_def` could be a dynamically-created message, so try a duck-typed
+    # approach
+    try:
+      old_graph_def = graph_def
+      graph_def = graph_pb2.GraphDef()
+      graph_def.MergeFrom(old_graph_def)
+    except TypeError:
+      raise TypeError('Argument `graph_def` must be a GraphDef proto.')
+  else:
+    # If we're using the graph_def provided by the caller, modify graph_def
+    # in-place to add attr defaults to the NodeDefs (this is visible to the
+    # caller).
+    # NOTE(skyewm): this is undocumented behavior that at least meta_graph.py
+    # depends on. It might make sense to move this to meta_graph.py and have
+    # import_graph_def not modify the graph_def argument (we'd have to make sure
+    # this doesn't break anything else.)
+    for node in graph_def.node:
+      op_def = op_def_registry.get(node.op)
+      if op_def is None:
+        # Assume unrecognized ops are functions for now. TF_ImportGraphDef will
+        # report an error if the op is actually missing.
+        continue
+      _SetDefaultAttrValues(node, op_def)
+
+  return graph_def
+
+
+def _ProcessInputMapParam(input_map):
+  """Type-checks and possibly canonicalizes `input_map`."""
+  if input_map is None:
+    input_map = {}
+  else:
+    if not isinstance(input_map, dict):
+      raise TypeError('Argument `input_map` must be a dictionary. Obtained '
+                      f'{type(input_map).__name__}')
+    if not all(
+        isinstance(k, compat.bytes_or_text_types) for k in input_map.keys()):
+      raise TypeError('All keys for argument `input_map` must be strings. '
+                      f'Obtained keys: {list(input_map.keys())}')
+  return input_map
+
+
+def _ProcessReturnElementsParam(return_elements):
+  """Type-checks and possibly canonicalizes `return_elements`."""
+  if return_elements is None:
+    return None
+  if not all(
+      isinstance(x, compat.bytes_or_text_types) for x in return_elements):
+    raise TypeError('Argument `return_elements` must be a list of strings. '
+                    f'Obtained {return_elements}.')
+  return tuple(compat.as_str(x) for x in return_elements)
+
+
+def _FindAttrInOpDef(attr_name, op_def):
+  for attr_def in op_def.attr:
+    if attr_name == attr_def.name:
+      return attr_def
+  return None
+
+
+def _RemoveDefaultAttrs(producer_op_list, graph_def):
+  """Removes unknown default attrs according to `producer_op_list`.
+
+  Removes any unknown attrs in `graph_def` (i.e. attrs that do not appear in
+  registered OpDefs) that have a default value in `producer_op_list`.
+
+  Args:
+    producer_op_list: OpList proto.
+    graph_def: GraphDef proto
+  """
+  producer_op_dict = {op.name: op for op in producer_op_list.op}
+  for node in graph_def.node:
+    # Remove any default attr values that aren't in op_def.
+    if node.op in producer_op_dict:
+      op_def = op_def_registry.get(node.op)
+      if op_def is None:
+        # Some custom op registrations won't show up here. That's OK, attribute
+        # stripping just won't be available.
+        continue
+      producer_op_def = producer_op_dict[node.op]
+      # We make a copy of node.attr to iterate through since we may modify
+      # node.attr inside the loop.
+      for key in list(node.attr):
+        if _FindAttrInOpDef(key, op_def) is None:
+          # No attr_def in consumer, look in producer.
+          attr_def = _FindAttrInOpDef(key, producer_op_def)
+          if (attr_def and attr_def.HasField('default_value') and
+              node.attr[key] == attr_def.default_value):
+            # Unknown attr had default value in producer, delete it so it can be
+            # understood by consumer.
+            del node.attr[key]
+
+
+def _ConvertInputMapValues(name, input_map):
+  """Ensures all input map values are tensors.
+
+  This should be called from inside the import name scope.
+
+  Args:
+    name: the `name` argument passed to import_graph_def
+    input_map: the `input_map` argument passed to import_graph_def.
+
+  Returns:
+    An possibly-updated version of `input_map`.
+
+  Raises:
+    ValueError: if input map values cannot be converted due to empty name scope.
+  """
+  if not all(isinstance(v, tensor.Tensor) for v in input_map.values()):
+    if name == '':  # pylint: disable=g-explicit-bool-comparison
+      raise ValueError(
+          'tf.import_graph_def() requires a non-empty `name` if `input_map` '
+          'contains non-Tensor values. Try calling tf.convert_to_tensor() on '
+          '`input_map` values before calling tf.import_graph_def().')
+    with ops.name_scope('_inputs'):
+      input_map = {k: ops.convert_to_tensor(v) for k, v in input_map.items()}
+  return input_map
+
+
+def _PopulateTFImportGraphDefOptions(options, prefix, input_map,
+                                     return_elements,
+                                     validate_colocation_constraints,
+                                     propagate_device_spec=False):
+  """Populates the TF_ImportGraphDefOptions `options`."""
+  c_api.TF_ImportGraphDefOptionsSetPrefix(options, prefix)
+  c_api.TF_ImportGraphDefOptionsSetUniquifyNames(options, True)
+  c_api.TF_ImportGraphDefOptionsSetPropagateDeviceSpec(options,
+                                                       propagate_device_spec)
+
+  for input_src, input_dst in input_map.items():
+    input_src = compat.as_str(input_src)
+    if input_src.startswith('^'):
+      src_name = compat.as_str(input_src[1:])
+      dst_op = input_dst._as_tf_output().oper  # pylint: disable=protected-access
+      c_api.TF_ImportGraphDefOptionsRemapControlDependency(
+          options, src_name, dst_op)
+    else:
+      src_name, src_idx = _ParseTensorName(input_src)
+      src_name = compat.as_str(src_name)
+      dst_output = input_dst._as_tf_output()  # pylint: disable=protected-access
+      c_api.TF_ImportGraphDefOptionsAddInputMapping(options, src_name, src_idx,
+                                                    dst_output)
+  for name in return_elements or []:
+    if ':' in name:
+      op_name, index = _ParseTensorName(name)
+      op_name = compat.as_str(op_name)
+      c_api.TF_ImportGraphDefOptionsAddReturnOutput(options, op_name, index)
+    else:
+      c_api.TF_ImportGraphDefOptionsAddReturnOperation(options,
+                                                       compat.as_str(name))
+
+  c_api.TF_ImportGraphDefOptionsSetValidateColocationConstraints(
+      options, validate_colocation_constraints)
+
+
+def _ProcessNewOps(graph):
+  """Processes the newly-added TF_Operations in `graph`."""
+  # Maps from a node to the names of the ops it's colocated with, if colocation
+  # is specified in the attributes.
+  colocation_pairs = {}
+
+  for new_op in graph._add_new_tf_operations(compute_devices=False):  # pylint: disable=protected-access
+    original_device = new_op.device
+    new_op._set_device('')  # pylint: disable=protected-access
+    colocation_names = _GetColocationNames(new_op)
+    if colocation_names:
+      colocation_pairs[new_op] = colocation_names
+      # Don't set a device for this op, since colocation constraints override
+      # device functions and the original device. Note that this op's device may
+      # still be set by the loop below.
+      # TODO(skyewm): why does it override the original device?
+    else:
+      with _MaybeDevice(original_device):
+        graph._apply_device_functions(new_op)  # pylint: disable=protected-access
+
+  # The following loop populates the device field of ops that are colocated
+  # with another op.  This is implied by the colocation attribute, but we
+  # propagate the device field for completeness.
+  for op, coloc_op_list in colocation_pairs.items():
+    coloc_device = None
+    # Find any device in the list of colocated ops that have a device, if it
+    # exists.  We assume that if multiple ops have devices, they refer to the
+    # same device.  Otherwise, a runtime error will occur since the colocation
+    # property cannot be guaranteed.  Note in TF2 colocations have been removed
+    # from the public API and will be considered a hint, so there is no runtime
+    # error.
+    #
+    # One possible improvement is to try to check for compatibility of all
+    # devices in this list at import time here, which would require
+    # implementing a compatibility function for device specs in python.
+    for coloc_op_name in coloc_op_list:
+      try:
+        coloc_op = graph._get_operation_by_name(coloc_op_name)  # pylint: disable=protected-access
+      except KeyError:
+        # Do not error in TF2 if the colocation cannot be guaranteed
+        if tf2.enabled() or control_flow_util.EnableControlFlowV2(graph):
+          continue
+
+        raise ValueError(f'Specified colocation to an op: {coloc_op_name} that '
+                         f'does not exist during import for op: {op.name}')
+      if coloc_op.device:
+        coloc_device = pydev.DeviceSpec.from_string(coloc_op.device)
+        break
+    if coloc_device:
+      op._set_device(coloc_device)  # pylint: disable=protected-access
+
+
+def _GetColocationNames(op):
+  """Returns names of the ops that `op` should be colocated with."""
+  colocation_names = []
+  try:
+    class_values = op.get_attr('_class')
+  except ValueError:
+    # No _class attr
+    return
+  for val in class_values:
+    val = compat.as_str(val)
+    if val.startswith('loc:@'):
+      colocation_node_name = val[len('loc:@'):]
+      if colocation_node_name != op.name:
+        colocation_names.append(colocation_node_name)
+  return colocation_names
+
+
+def _GatherReturnElements(requested_return_elements, graph, results):
+  """Returns the requested return elements from results.
+
+  Args:
+    requested_return_elements: list of strings of operation and tensor names
+    graph: Graph
+    results: wrapped TF_ImportGraphDefResults
+
+  Returns:
+    list of `Operation` and/or `Tensor` objects
+  """
+  return_outputs = c_api.TF_ImportGraphDefResultsReturnOutputs(results)
+  return_opers = c_api.TF_ImportGraphDefResultsReturnOperations(results)
+
+  combined_return_elements = []
+  outputs_idx = 0
+  opers_idx = 0
+  for name in requested_return_elements:
+    if ':' in name:
+      combined_return_elements.append(
+          graph._get_tensor_by_tf_output(return_outputs[outputs_idx]))  # pylint: disable=protected-access
+      outputs_idx += 1
+    else:
+      combined_return_elements.append(
+          graph._get_operation_by_tf_operation(return_opers[opers_idx]))  # pylint: disable=protected-access
+      opers_idx += 1
+  return combined_return_elements
+
+
+def _SetDefaultAttrValues(node_def, op_def):
+  """Set any default attr values in `node_def` that aren't present."""
+  assert node_def.op == op_def.name
+  for attr_def in op_def.attr:
+    key = attr_def.name
+    if attr_def.HasField('default_value'):
+      value = node_def.attr[key]
+      if value is None or value.WhichOneof('value') is None:
+        node_def.attr[key].CopyFrom(attr_def.default_value)
+
+
+@tf_export('graph_util.import_graph_def', 'import_graph_def')
+@deprecated_args(None, 'Please file an issue at '
+                 'https://github.com/tensorflow/tensorflow/issues if you depend'
+                 ' on this feature.', 'op_dict')
+def import_graph_def(graph_def,
+                     input_map=None,
+                     return_elements=None,
+                     name=None,
+                     op_dict=None,
+                     producer_op_list=None):
+  """Imports the graph from `graph_def` into the current default `Graph`.
+
+  This function provides a way to import a serialized TensorFlow
+  [`GraphDef`](https://www.tensorflow.org/code/tensorflow/core/framework/graph.proto)
+  protocol buffer, and extract individual objects in the `GraphDef` as
+  `tf.Tensor` and `tf.Operation` objects. Once extracted,
+  these objects are placed into the current default `Graph`. See
+  `tf.Graph.as_graph_def` for a way to create a `GraphDef`
+  proto.
+
+  Args:
+    graph_def: A `GraphDef` proto containing operations to be imported into
+      the default graph.
+    input_map: A dictionary mapping input names (as strings) in `graph_def`
+      to `Tensor` objects. The values of the named input tensors in the
+      imported graph will be re-mapped to the respective `Tensor` values.
+    return_elements: A list of strings containing operation names in
+      `graph_def` that will be returned as `Operation` objects; and/or
+      tensor names in `graph_def` that will be returned as `Tensor` objects.
+    name: (Optional.) A prefix that will be prepended to the names in
+      `graph_def`. Note that this does not apply to imported function names.
+      Defaults to `"import"`.
+    op_dict: (Optional.) Deprecated, do not use.
+    producer_op_list: (Optional.) An `OpList` proto with the (possibly stripped)
+      list of `OpDef`s used by the producer of the graph. If provided,
+      unrecognized attrs for ops in `graph_def` that have their default value
+      according to `producer_op_list` will be removed. This will allow some more
+      `GraphDef`s produced by later binaries to be accepted by earlier binaries.
+
+  Returns:
+    A list of `Operation` and/or `Tensor` objects from the imported graph,
+    corresponding to the names in `return_elements`,
+    and None if `returns_elements` is None.
+
+  Raises:
+    TypeError: If `graph_def` is not a `GraphDef` proto,
+      `input_map` is not a dictionary mapping strings to `Tensor` objects,
+      or `return_elements` is not a list of strings.
+    ValueError: If `input_map`, or `return_elements` contains names that
+      do not appear in `graph_def`, or `graph_def` is not well-formed (e.g.
+      it refers to an unknown tensor).
+  """
+  del op_dict
+  return _import_graph_def_internal(
+      graph_def,
+      input_map=input_map,
+      return_elements=return_elements,
+      name=name,
+      producer_op_list=producer_op_list)
+
+
+def import_graph_def_for_function(  # pylint: disable=invalid-name
+    graph_def, name=None, propagate_device_spec=False):
+  """Like import_graph_def but does not validate colocation constraints."""
+  return _import_graph_def_internal(
+      graph_def,
+      validate_colocation_constraints=False,
+      name=name,
+      propagate_device_spec=propagate_device_spec)
+
+
+def _import_graph_def_internal(  # pylint: disable=invalid-name
+    graph_def,
+    input_map=None,
+    return_elements=None,
+    validate_colocation_constraints=True,
+    name=None,
+    producer_op_list=None,
+    propagate_device_spec=False):
+  """Imports the graph from `graph_def` into the current default `Graph`.
+
+  This function provides a way to import a serialized TensorFlow
+  [`GraphDef`](https://www.tensorflow.org/code/tensorflow/core/framework/graph.proto)
+  protocol buffer, and extract individual objects in the `GraphDef` as
+  `tf.Tensor` and `tf.Operation` objects. Once extracted,
+  these objects are placed into the current default `Graph`. See
+  `tf.Graph.as_graph_def` for a way to create a `GraphDef`
+  proto.
+
+  Args:
+    graph_def: A `GraphDef` proto containing operations to be imported into the
+      default graph.
+    input_map: A dictionary mapping input names (as strings) in `graph_def` to
+      `Tensor` objects. The values of the named input tensors in the imported
+      graph will be re-mapped to the respective `Tensor` values.
+    return_elements: A list of strings containing operation names in `graph_def`
+      that will be returned as `Operation` objects; and/or tensor names in
+      `graph_def` that will be returned as `Tensor` objects.
+    validate_colocation_constraints: Whether to validate colocation constraints.
+    name: (Optional.) A prefix that will be prepended to the names in
+      `graph_def`. Note that this does not apply to imported function names.
+      Defaults to `"import"`.
+    producer_op_list: (Optional.) An `OpList` proto with the (possibly stripped)
+      list of `OpDef`s used by the producer of the graph. If provided,
+      unrecognized attrs for ops in `graph_def` that have their default value
+      according to `producer_op_list` will be removed. This will allow some more
+      `GraphDef`s produced by later binaries to be accepted by earlier binaries.
+    propagate_device_spec: Whether to propagate assigned device information
+      when importing a graph from a GraphDef into the current default `Graph`.
+
+  Returns:
+    A list of `Operation` and/or `Tensor` objects from the imported graph,
+    corresponding to the names in `return_elements`,
+    and None if `returns_elements` is None.
+
+  Raises:
+    TypeError: If `graph_def` is not a `GraphDef` proto,
+      `input_map` is not a dictionary mapping strings to `Tensor` objects,
+      or `return_elements` is not a list of strings.
+    ValueError: If `input_map`, or `return_elements` contains names that
+      do not appear in `graph_def`, or `graph_def` is not well-formed (e.g.
+      it refers to an unknown tensor).
+  """
+  graph_def = _ProcessGraphDefParam(graph_def)
+  input_map = _ProcessInputMapParam(input_map)
+  return_elements = _ProcessReturnElementsParam(return_elements)
+
+  if producer_op_list is not None:
+    # TODO(skyewm): make a copy of graph_def so we're not mutating the argument?
+    _RemoveDefaultAttrs(producer_op_list, graph_def)
+
+  graph = ops.get_default_graph()
+  with ops.name_scope(name, 'import', input_map.values()) as scope:
+    # Save unique prefix generated by name_scope
+    if scope:
+      assert scope.endswith('/')
+      prefix = scope[:-1]
+    else:
+      prefix = ''
+
+    # Generate any input map tensors inside name scope
+    input_map = _ConvertInputMapValues(name, input_map)
+
+  scoped_options = c_api_util.ScopedTFImportGraphDefOptions()
+  options = scoped_options.options
+  _PopulateTFImportGraphDefOptions(options, prefix, input_map, return_elements,
+                                   validate_colocation_constraints,
+                                   propagate_device_spec)
+
+  # _ProcessNewOps mutates the new operations. _mutation_lock ensures a
+  # Session.run call cannot occur between creating the TF_Operations in the
+  # TF_GraphImportGraphDefWithResults call and mutating them in _ProcessNewOps.
+  with graph._mutation_lock():  # pylint: disable=protected-access
+    if is_oss:
+      graph_def_input = c_api.TF_NewBufferFromString(
+          compat.as_bytes(graph_def.SerializeToString())
+      )
+      graph_import_graphdef = c_api.TF_GraphImportGraphDefWithResults
+    else:
+      graph_def_input = graph_def
+      graph_import_graphdef = (
+          c_api.TF_GraphImportGraphDefWithResultsNoSerialization
+      )
+    try:
+      with graph._c_graph.get() as c_graph:  # pylint: disable=protected-access
+        results = graph_import_graphdef(c_graph, graph_def_input, options)
+      results = c_api_util.ScopedTFImportGraphDefResults(results)
+    except errors.InvalidArgumentError as e:
+      # Convert to ValueError for backwards compatibility.
+      raise ValueError(str(e))
+    finally:
+      if is_oss:
+        c_api.TF_DeleteBuffer(graph_def_input)
+
+    # Create _DefinedFunctions for any imported functions.
+    #
+    # We do this by creating _DefinedFunctions directly from `graph_def`, and
+    # adding them to `graph`. Adding an existing function to a TF_Graph is a
+    # no-op, so this only has the effect of updating the Python state (usually
+    # _DefinedFunction.add_to_graph also adds the function to the TF_Graph).
+    #
+    # TODO(skyewm): fetch the TF_Functions directly from the TF_Graph
+    # TODO(skyewm): avoid sending serialized FunctionDefs back to the TF_Graph
+
+    _ProcessNewOps(graph)
+
+  if graph_def.library and graph_def.library.function:
+    functions = function.from_library(graph_def.library)
+    for f in functions:
+      f.add_to_graph(graph)
+
+  # Treat input mappings that don't appear in the graph as an error, because
+  # they are likely to be due to a typo.
+  missing_unused_input_keys = (
+      c_api.TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper(
+          results.results))
+  if missing_unused_input_keys:
+    missing_unused_input_keys = [
+        compat.as_str(s) for s in missing_unused_input_keys
+    ]
+    missing_keys = ', '.join(missing_unused_input_keys)
+    raise ValueError(
+        'Attempted to map inputs that were not found in graph_def: '
+        f'[{missing_keys}]')
+
+  if return_elements is None:
+    return None
+  else:
+    return _GatherReturnElements(return_elements, graph, results.results)

tensorflow/python/framework/indexed_slices.py

@@ -0,0 +1,455 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Indexed slices."""
+
+# pylint: disable=g-bad-name
+import collections
+import warnings
+
+import numpy as np
+
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python import tf2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import composite_tensor_gradient
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion_registry
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.types import internal
+from tensorflow.python.util.compat import collections_abc
+from tensorflow.python.util.tf_export import tf_export
+
+
+class IndexedSlicesCompositeTensorGradient(
+    composite_tensor_gradient.CompositeTensorGradient):
+  """CompositeTensorGradient for IndexedSlices."""
+
+  def get_gradient_components(self, value):
+    return value
+
+  def replace_gradient_components(self, value, component_grads):
+    return component_grads
+
+
+# TODO(mdan): Should IndexedSlices be a "tensor"?
+@tf_export("IndexedSlices")
+class IndexedSlices(
+    internal.IndexedSlices,
+    internal.NativeObject,
+    composite_tensor.CompositeTensor):
+  """A sparse representation of a set of tensor slices at given indices.
+
+  This class is a simple wrapper for a pair of `Tensor` objects:
+
+  * `values`: A `Tensor` of any dtype with shape `[D0, D1, ..., Dn]`.
+  * `indices`: A 1-D integer `Tensor` with shape `[D0]`.
+
+  An `IndexedSlices` is typically used to represent a subset of a larger
+  tensor `dense` of shape `[LARGE0, D1, .. , DN]` where `LARGE0 >> D0`.
+  The values in `indices` are the indices in the first dimension of
+  the slices that have been extracted from the larger tensor.
+
+  The dense tensor `dense` represented by an `IndexedSlices` `slices` has
+
+  ```python
+  dense[slices.indices[i], :, :, :, ...] = slices.values[i, :, :, :, ...]
+  ```
+
+  The `IndexedSlices` class is used principally in the definition of
+  gradients for operations that have sparse gradients
+  (e.g. `tf.gather`).
+
+  >>> v = tf.Variable([[0.,1, 2], [2, 3, 4], [4, 5, 6], [6, 7, 8]])
+  >>> with tf.GradientTape() as tape:
+  ...   r = tf.gather(v, [1,3])
+  >>> index_slices = tape.gradient(r,v)
+  >>> index_slices
+  <...IndexedSlices object ...>
+  >>> index_slices.indices.numpy()
+  array([1, 3], dtype=int32)
+  >>> index_slices.values.numpy()
+  array([[1., 1., 1.],
+         [1., 1., 1.]], dtype=float32)
+
+  Contrast this representation with
+  `tf.sparse.SparseTensor`,
+  which uses multi-dimensional indices and scalar values.
+  """
+
+  def __init__(self, values, indices, dense_shape=None):
+    """Creates an `IndexedSlices`."""
+    self._values = values
+    self._indices = indices
+    self._dense_shape = dense_shape
+
+  @property
+  def values(self):
+    """A `Tensor` containing the values of the slices."""
+    return self._values
+
+  @property
+  def indices(self):
+    """A 1-D `Tensor` containing the indices of the slices."""
+    return self._indices
+
+  @property
+  def dense_shape(self):
+    """A 1-D `Tensor` containing the shape of the corresponding dense tensor."""
+    return self._dense_shape
+
+  @property
+  def shape(self):
+    """Gets the `tf.TensorShape` representing the shape of the dense tensor.
+
+    Returns:
+      A `tf.TensorShape` object.
+    """
+    if self._dense_shape is None:
+      return tensor_shape.TensorShape(None)
+
+    return tensor_util.constant_value_as_shape(self._dense_shape)
+
+  @property
+  def name(self):
+    """The name of this `IndexedSlices`."""
+    return self.values.name
+
+  @property
+  def device(self):
+    """The name of the device on which `values` will be produced, or `None`."""
+    return self.values.device
+
+  @property
+  def op(self) -> ops.Operation:
+    """The `Operation` that produces `values` as an output."""
+    return self.values.op
+
+  @property
+  def dtype(self):
+    """The `DType` of elements in this tensor."""
+    return self.values.dtype
+
+  @property
+  def graph(self) -> ops.Graph:
+    """The `Graph` that contains the values, indices, and shape tensors."""
+    return self._values.graph
+
+  def __str__(self):
+    return "IndexedSlices(indices=%s, values=%s%s)" % (
+        self._indices, self._values,
+        (", dense_shape=%s" %
+         (self._dense_shape,)) if self._dense_shape is not None else "")
+
+  def __neg__(self):
+    return IndexedSlices(-self.values, self.indices, self.dense_shape)
+
+  __composite_gradient__ = IndexedSlicesCompositeTensorGradient()
+
+  @property
+  def _type_spec(self):
+    indices_shape = self._indices.shape.merge_with(self._values.shape[:1])
+    dense_shape = tensor_shape.TensorShape([None]).concatenate(
+        self._values.shape[1:])
+    if self._dense_shape is not None:
+      dense_shape_dtype = self._dense_shape.dtype
+      dense_shape = dense_shape.merge_with(
+          tensor_util.constant_value_as_shape(self._dense_shape))
+    else:
+      dense_shape_dtype = None
+    return IndexedSlicesSpec(dense_shape, self.dtype, self._indices.dtype,
+                             dense_shape_dtype, indices_shape)
+
+  def _shape_invariant_to_type_spec(self, shape):
+    # From tf.while_loop docs: "If a loop variable is an IndexedSlices, the
+    # shape invariant must be a shape invariant of the values tensor of the
+    # IndexedSlices. It means the shapes of the three tensors of the
+    # IndexedSlices are (shape, [shape[0]], [shape.ndims])."
+    indices_shape = shape[:1]
+    dense_shape = tensor_shape.TensorShape([None]).concatenate(shape[1:])
+    if self._dense_shape is None:
+      dense_shape_dtype = None
+    else:
+      dense_shape_dtype = self._dense_shape.dtype
+    return IndexedSlicesSpec(dense_shape, self.dtype, self._indices.dtype,
+                             dense_shape_dtype, indices_shape)
+
+  def consumers(self):
+    return self._consumers()
+
+
+IndexedSlicesValue = collections.namedtuple(
+    "IndexedSlicesValue", ["values", "indices", "dense_shape"])
+
+
+@tf_export("IndexedSlicesSpec")
+class IndexedSlicesSpec(type_spec.TypeSpec):
+  """Type specification for a `tf.IndexedSlices`."""
+
+  __slots__ = ["_shape", "_values_dtype", "_indices_dtype",
+               "_dense_shape_dtype", "_indices_shape"]
+
+  value_type = property(lambda self: IndexedSlices)
+
+  def __init__(self, shape=None, dtype=dtypes.float32,
+               indices_dtype=dtypes.int64, dense_shape_dtype=None,
+               indices_shape=None):
+    """Constructs a type specification for a `tf.IndexedSlices`.
+
+    Args:
+      shape: The dense shape of the `IndexedSlices`, or `None` to allow any
+        dense shape.
+      dtype: `tf.DType` of values in the `IndexedSlices`.
+      indices_dtype: `tf.DType` of the `indices` in the `IndexedSlices`.  One
+        of `tf.int32` or `tf.int64`.
+      dense_shape_dtype: `tf.DType` of the `dense_shape` in the `IndexedSlices`.
+        One of `tf.int32`, `tf.int64`, or `None` (if the `IndexedSlices` has
+        no `dense_shape` tensor).
+      indices_shape: The shape of the `indices` component, which indicates
+        how many slices are in the `IndexedSlices`.
+    """
+    self._shape = tensor_shape.as_shape(shape)
+    self._values_dtype = dtypes.as_dtype(dtype)
+    self._indices_dtype = dtypes.as_dtype(indices_dtype)
+    if dense_shape_dtype is None:
+      self._dense_shape_dtype = None
+    else:
+      self._dense_shape_dtype = dtypes.as_dtype(dense_shape_dtype)
+    self._indices_shape = tensor_shape.as_shape(indices_shape).with_rank(1)
+
+  def _serialize(self):
+    return (self._shape, self._values_dtype, self._indices_dtype,
+            self._dense_shape_dtype, self._indices_shape)
+
+  @property
+  def _component_specs(self):
+    value_shape = self._indices_shape.concatenate(self._shape[1:])
+    specs = [
+        tensor_spec.TensorSpec(value_shape, self._values_dtype),
+        tensor_spec.TensorSpec(self._indices_shape, self._indices_dtype)]
+    if self._dense_shape_dtype is not None:
+      specs.append(
+          tensor_spec.TensorSpec([self._shape.ndims], self._dense_shape_dtype))
+    return tuple(specs)
+
+  def _to_components(self, value):
+    if value.dense_shape is None:
+      return (value.values, value.indices)
+    else:
+      return (value.values, value.indices, value.dense_shape)
+
+  def _from_components(self, tensor_list):
+    if (all(isinstance(t, np.ndarray) for t in tensor_list) and
+        not tf2.enabled()):
+      if len(tensor_list) == 2:
+        return IndexedSlicesValue(tensor_list[0], tensor_list[1], None)
+      else:
+        return IndexedSlicesValue(*tensor_list)
+    else:
+      return IndexedSlices(*tensor_list)
+
+
+nested_structure_coder.register_codec(
+    nested_structure_coder.BuiltInTypeSpecCodec(
+        IndexedSlicesSpec, struct_pb2.TypeSpecProto.INDEXED_SLICES_SPEC
+    )
+)
+
+
+@tf_export(v1=["convert_to_tensor_or_indexed_slices"])
+def convert_to_tensor_or_indexed_slices(value, dtype=None, name=None):
+  """Converts the given object to a `Tensor` or an `IndexedSlices`.
+
+  If `value` is an `IndexedSlices` or `SparseTensor` it is returned
+  unmodified. Otherwise, it is converted to a `Tensor` using
+  `convert_to_tensor()`.
+
+  Args:
+    value: An `IndexedSlices`, `SparseTensor`, or an object that can be consumed
+      by `convert_to_tensor()`.
+    dtype: (Optional.) The required `DType` of the returned `Tensor` or
+      `IndexedSlices`.
+    name: (Optional.) A name to use if a new `Tensor` is created.
+
+  Returns:
+    A `Tensor`, `IndexedSlices`, or `SparseTensor` based on `value`.
+
+  Raises:
+    ValueError: If `dtype` does not match the element type of `value`.
+  """
+  return internal_convert_to_tensor_or_indexed_slices(
+      value=value, dtype=dtype, name=name, as_ref=False)
+
+
+def internal_convert_to_tensor_or_indexed_slices(value,
+                                                 dtype=None,
+                                                 name=None,
+                                                 as_ref=False):
+  """Converts the given object to a `Tensor` or an `IndexedSlices`.
+
+  If `value` is an `IndexedSlices` or `SparseTensor` it is returned
+  unmodified. Otherwise, it is converted to a `Tensor` using
+  `convert_to_tensor()`.
+
+  Args:
+    value: An `IndexedSlices`, `SparseTensor`, or an object that can be consumed
+      by `convert_to_tensor()`.
+    dtype: (Optional.) The required `DType` of the returned `Tensor` or
+      `IndexedSlices`.
+    name: (Optional.) A name to use if a new `Tensor` is created.
+    as_ref: True if the caller wants the results as ref tensors.
+
+  Returns:
+    A `Tensor`, `IndexedSlices`, or `SparseTensor` based on `value`.
+
+  Raises:
+    ValueError: If `dtype` does not match the element type of `value`.
+  """
+  if isinstance(value, ops.EagerTensor) and not context.executing_eagerly():
+    return ops.convert_to_tensor(value, dtype=dtype, name=name, as_ref=as_ref)
+  # TODO(mdan): Name says tensor_or_indexed_slices. So do explicitly just that?
+  elif isinstance(value, internal.NativeObject):
+    if dtype and not dtypes.as_dtype(dtype).is_compatible_with(value.dtype):
+      raise ValueError(
+          "Incompatible tensor conversion requested to `dtype` "
+          f"{dtypes.as_dtype(dtype).name} for `value` ({value}) with dtype"
+          f" {value.dtype.name}.")
+    return value
+  else:
+    return ops.convert_to_tensor(value, dtype=dtype, name=name, as_ref=as_ref)
+
+
+def internal_convert_n_to_tensor_or_indexed_slices(values,
+                                                   dtype=None,
+                                                   name=None,
+                                                   as_ref=False):
+  """Converts `values` to a list of `Tensor` or `IndexedSlices` objects.
+
+  Any `IndexedSlices` or `SparseTensor` objects in `values` are returned
+  unmodified.
+
+  Args:
+    values: An iterable of `None`, `IndexedSlices`, `SparseTensor`, or objects
+      that can be consumed by `convert_to_tensor()`.
+    dtype: (Optional.) The required `DType` of the returned `Tensor` or
+      `IndexedSlices`.
+    name: (Optional.) A name prefix to used when a new `Tensor` is created, in
+      which case element `i` will be given the name `name + '_' + i`.
+    as_ref: True if the caller wants the results as ref tensors.
+
+  Returns:
+    A list of `Tensor`, `IndexedSlices`, `SparseTensor` and/or `None` objects.
+
+  Raises:
+    TypeError: If no conversion function is registered for an element in
+      `values`.
+    RuntimeError: If a registered conversion function returns an invalid
+      value.
+  """
+  if not isinstance(values, collections_abc.Iterable):
+    raise TypeError("Argument `values` must be iterable.")
+  ret = []
+  for i, value in enumerate(values):
+    if value is None:
+      ret.append(value)
+    else:
+      n = None if name is None else "%s_%d" % (name, i)
+      ret.append(
+          internal_convert_to_tensor_or_indexed_slices(
+              value, dtype=dtype, name=n, as_ref=as_ref))
+  return ret
+
+
+def convert_n_to_tensor_or_indexed_slices(values, dtype=None, name=None):
+  """Converts `values` to a list of `Output` or `IndexedSlices` objects.
+
+  Any `IndexedSlices` or `SparseTensor` objects in `values` are returned
+  unmodified.
+
+  Args:
+    values: A list of `None`, `IndexedSlices`, `SparseTensor`, or objects that
+      can be consumed by `convert_to_tensor()`.
+    dtype: (Optional.) The required `DType` of the returned `Tensor`
+      `IndexedSlices`.
+    name: (Optional.) A name prefix to used when a new `Tensor` is created, in
+      which case element `i` will be given the name `name + '_' + i`.
+
+  Returns:
+    A list of `Tensor`, `IndexedSlices`, and/or `SparseTensor` objects.
+
+  Raises:
+    TypeError: If no conversion function is registered for an element in
+      `values`.
+    RuntimeError: If a registered conversion function returns an invalid
+      value.
+  """
+  return internal_convert_n_to_tensor_or_indexed_slices(
+      values=values, dtype=dtype, name=name, as_ref=False)
+
+
+# Warn the user if we convert a sparse representation to dense with at
+# least this number of elements.
+_LARGE_SPARSE_NUM_ELEMENTS = 100000000
+
+
+def _indexed_slices_to_tensor(value, dtype=None, name=None, as_ref=False):
+  """Converts an IndexedSlices object `value` to a Tensor.
+
+  NOTE(mrry): This function is potentially expensive.
+
+  Args:
+    value: An ops.IndexedSlices object.
+    dtype: The dtype of the Tensor to be returned.
+    name: Optional name to use for the returned Tensor.
+    as_ref: True if a ref is requested.
+
+  Returns:
+    A dense Tensor representing the values in the given IndexedSlices.
+
+  Raises:
+    ValueError: If the IndexedSlices does not have the same dtype.
+  """
+  _ = as_ref
+  if dtype and not dtype.is_compatible_with(value.dtype):
+    raise ValueError(
+        f"Incompatible tensor conversion requested to `dtype` {dtype.name} for "
+        f"IndexedSlices ({value}) with dtype {value.dtype.name}")
+  if value.dense_shape is None:
+    raise ValueError(
+        "Tensor conversion requested for IndexedSlices for argument `value` "
+        f"without dense_shape: {value!s}")
+  # TODO(mrry): Consider adding static shape information to
+  # IndexedSlices, to avoid using numpy here.
+  if not context.executing_eagerly():
+    dense_shape_value = tensor_util.constant_value(value.dense_shape)
+    if dense_shape_value is not None:
+      num_elements = np.prod(dense_shape_value)
+      if num_elements >= _LARGE_SPARSE_NUM_ELEMENTS:
+        warnings.warn(
+            "Converting sparse IndexedSlices to a dense Tensor with %d "
+            "elements. This may consume a large amount of memory." %
+            num_elements)
+  return gen_math_ops.unsorted_segment_sum(
+      value.values, value.indices, value.dense_shape[0], name=name)
+
+
+tensor_conversion_registry.register_tensor_conversion_function(
+    IndexedSlices, _indexed_slices_to_tensor)

tensorflow/python/framework/indexed_slices_test.py

@@ -0,0 +1,40 @@
+# 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 third_party.tensorflow.python.framework.indexed_slices."""
+
+from tensorflow.python.framework import composite_tensor_gradient
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.platform import test
+
+
+class IndexedSlicesTest(test.TestCase):
+
+  def testCompositeTensorGradient(self):
+    i = indexed_slices.IndexedSlices(values=constant_op.constant([[1., 2.]]),
+                                     indices=constant_op.constant([1]),
+                                     dense_shape=[3, 2])
+    gradient_components = (
+        composite_tensor_gradient.get_flat_tensors_for_gradients([i]))
+    self.assertAllEqual(gradient_components, [i])
+
+    t = [3., 4.]
+    result = (
+        composite_tensor_gradient.replace_flat_tensors_for_gradients([i], [t]))
+    self.assertAllEqual(result, [t])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/framework/is_mlir_bridge_test_false.py

@@ -0,0 +1,26 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Including this as a dependency will result in tests NOT using MLIR bridge.
+
+This function is defined by default in test_util.py to None. The test_util then
+attempts to import this module. If this file is made available through the BUILD
+rule, then this function is overridden and will instead cause Tensorflow graphs
+to be always NOT be compiled with MLIR bridge.
+"""
+
+
+def is_mlir_bridge_enabled():
+  """Returns false if the MLIR bridge should be not be enabled for tests."""
+  return False

tensorflow/python/framework/is_mlir_bridge_test_true.py

@@ -0,0 +1,26 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Including this as a dependency will result in tests using MLIR bridge.
+
+This function is defined by default in test_util.py to False. The test_util then
+attempts to import this module. If this file is made available through the BUILD
+rule, then this function is overridden and will instead cause Tensorflow graphs
+to be compiled with MLIR bridge.
+"""
+
+
+def is_mlir_bridge_enabled():
+  """Returns true if MLIR bridge should be enabled for tests."""
+  return True

tensorflow/python/framework/is_xla_test_true.py

@@ -0,0 +1,25 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Including this as a dependency will result in Tensorflow tests using XLA.
+
+This function is defined by default in test_util.py to False. The test_util then
+attempts to import this module. If this file is made available through the BUILD
+rule, then this function is overridden and will instead cause Tensorflow graphs
+to be compiled with XLA.
+"""
+
+def is_xla_enabled():
+  """Returns true to state XLA should be enabled for Tensorflow tests."""
+  return True

tensorflow/python/framework/kernels.py

@@ -0,0 +1,42 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Functions for querying registered kernels."""
+
+from tensorflow.core.framework import kernel_def_pb2
+from tensorflow.python.client import pywrap_tf_session as c_api
+from tensorflow.python.util import compat
+
+
+def get_all_registered_kernels():
+  """Returns a KernelList proto of all registered kernels.
+  """
+  buf = c_api.TF_GetAllRegisteredKernels()
+  data = c_api.TF_GetBuffer(buf)
+  kernel_list = kernel_def_pb2.KernelList()
+  kernel_list.ParseFromString(compat.as_bytes(data))
+  return kernel_list
+
+
+def get_registered_kernels_for_op(name):
+  """Returns a KernelList proto of registered kernels for a given op.
+
+  Args:
+    name: A string representing the name of the op whose kernels to retrieve.
+  """
+  buf = c_api.TF_GetRegisteredKernelsForOp(name)
+  data = c_api.TF_GetBuffer(buf)
+  kernel_list = kernel_def_pb2.KernelList()
+  kernel_list.ParseFromString(compat.as_bytes(data))
+  return kernel_list

tensorflow/python/framework/kernels_test.py

@@ -0,0 +1,37 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for querying registered kernels."""
+from tensorflow.python.framework import kernels
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+
+
+class GetAllRegisteredKernelsTest(test_util.TensorFlowTestCase):
+
+  def testFindsAtLeastOneKernel(self):
+    kernel_list = kernels.get_all_registered_kernels()
+    self.assertGreater(len(kernel_list.kernel), 0)
+
+
+class GetRegisteredKernelsForOp(test_util.TensorFlowTestCase):
+
+  def testFindsAtLeastOneKernel(self):
+    kernel_list = kernels.get_registered_kernels_for_op("KernelLabel")
+    self.assertGreater(len(kernel_list.kernel), 0)
+    self.assertEqual(kernel_list.kernel[0].op, "KernelLabel")
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/framework/kythe_metadata.proto

@@ -0,0 +1,142 @@
+/*
+ * Copyright 2019 The Kythe 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.
+ */
+
+/*
+ * The content of this file is copied from
+ * https://github.com/kythe/kythe/blob/master/kythe/proto/metadata.proto and
+ * https://github.com/kythe/kythe/blob/master/kythe/proto/storage.proto
+ * and shouldn't be modified unless the protos is outdated.
+ */
+
+syntax = "proto3";
+
+package tensorflow;
+
+// Schema for the JSON-encoded Kythe metadata describing the relationship
+// between source and target code for generated code.
+message GeneratedCodeInfo {
+  enum Type {
+    NONE = 0;
+    KYTHE0 = 1;  // Initial metadata document type.
+  }
+
+  Type type = 1;
+  repeated MappingRule meta = 2;  // Only relevant if type == kythe0.
+}
+
+// Metadata for a single mapping between a generated source range and a node
+// in the source language or file.
+message MappingRule {
+  enum Type {
+    NONE = 0;
+    NOP = 1;             // Dummy rule that contains no relevant information.
+    ANCHOR_DEFINES = 2;  // Rule describing a generates edge between target
+                         // range and source definition.
+    ANCHOR_ANCHOR = 3;   // Rule describing an imputes edge between target range
+                         // and source range.
+  }
+
+  Type type = 1;
+  // If type == anchor_defines, this should generally be a reverse generates
+  // edge, %/kythe/edge/generates, indicating that the specified vname generated
+  // the source range.
+  // If type == anchor_anchor, this should generally be a forward imputes edge,
+  // /kythe/edge/imputes, indicating that the range in the source file produced
+  // the text in the target file.
+  // If semantic is not NONE, this field is ignored and the identified
+  // declaration at the indicated text range is given the associated semantic.
+  string edge = 2;
+
+  // Fields only relevant if type == anchor_defines.
+  VName vname = 3;   // The semantic node in the source language
+                     // which generated the text range.
+  uint32 begin = 4;  // Beginning of the range to match in the generated text.
+  uint32 end = 5;    // End of the range to match in the generated text.
+
+  enum Semantic {
+    SEMA_NONE = 0;
+    SEMA_WRITE = 1;
+    SEMA_READ_WRITE = 2;
+  }
+
+  Semantic semantic = 11;
+
+  // Fields only relevant if type == anchor_anchor.
+  VName source_vname = 6;   // Anchor node in the generating source file.
+                            // Note: the signature in this vname, if present,
+                            // will typically be replaced by the target indexer
+                            // using its own anchor-construction rules based on
+                            // source_begin and source_end.
+  uint32 source_begin = 7;  // loc/start of the anchor node in the source file.
+  uint32 source_end = 8;    // loc/end of the anchor node in the source file.
+  uint32 target_begin = 9;  // Start of the range in the generated text.
+  uint32 target_end = 10;   // End of the range in the generated text.
+}
+
+message VName {
+  // A language-specific signature assigned by the analyzer.
+  // e.g., "com.google.common.collect.Lists.newLinkedList<#1>()"
+  string signature = 1;
+
+  // The corpus this name belongs to.
+  // e.g., "kythe", "chromium", "github.com/creachadair/imath", "aosp"
+  // The corpus label "kythe" is reserved for internal use.
+  string corpus = 2;
+
+  // A corpus-specific root label, designating a subordinate collection within
+  // the corpus.  If a corpus stores files in unrelated directory structures,
+  // for example, the root can be used to distinguish them.  Or, if a corpus
+  // incorporates subprojects, the root can be a project ID that it governs.
+  // This may also be used to distinguish virtual subgroups of a corpus such as
+  // generated files.
+  string root = 3;
+
+  // A path-structured label describing the location of this object relative to
+  // the corpus and the root.  For code, this will generally be the relative
+  // path to the file containing the code, e.g., "storage/service.go" in kythe.
+  // The individual elements of the path are separated by "/".
+  // The path must not start with "/".
+  // The path must be normalized to a canonical form (with no path
+  // elements "", ".", or "..").
+  //
+  // However, this need not be a true file path; virtual objects like figments
+  // can assign an ad-hoc abstract ID, or omit it entirely.
+  //
+  // Examples:
+  //   "devools/kythe/platform/go/datastore.go" (a file)
+  //   "type/cpp/void.cc" (a type figment)
+  string path = 4;
+
+  // The language this name belongs to.
+  // e.g., "c++", "python", "elisp", "haskell", "java"
+  //
+  // The schema will define specific labels for each supported language, so we
+  // don't wind up with a confusion of names like "cxx", "cpp", "C++", etc.
+  // Prototype: Official language name converted to lowercase.  If a version
+  // number is necessary, include it, e.g., "python3".
+  string language = 5;
+
+  // Other fields we may need in the future, but do not currently use:
+  // branch -- a branch name within the corpus depot, e.g., "gslb_branch".
+  // client -- a source-control client ID, e.g., "sergey:googlex:8:citc".
+
+  // Note: We have intentionally NOT included a revision or timestamp here.
+  // Time should be recorded as facts belonging to the appropriate Nodes and
+  // Edges.  Having records of when something existed may be important, but time
+  // is not a good axis for a name -- a name should say "what" something is, not
+  // "when".  So we will store timestamps, revisions, and other markers of this
+  // kind as facts inside the graph.
+}