chore: import upstream snapshot with attribution

tensorflow/python/util/_function_parameter_canonicalizer_binding_for_test.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 FunctionParameterCanonicalizer:
+    def __init__(self, arg0: list[str], arg1: tuple) -> None: ...
+    def canonicalize(self, *args, **kwargs) -> list: ...

tensorflow/python/util/_pywrap_checkpoint_reader.pyi

@@ -0,0 +1,21 @@
+# 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 CheckpointReader:
+    def __init__(self, arg0: str) -> None: ...
+    @staticmethod
+    def CheckpointReader_GetTensor(arg0: CheckpointReader, arg1: str) -> object: ...
+    def debug_string(self) -> bytes: ...
+    def get_variable_to_shape_map(self, *args, **kwargs): ...

tensorflow/python/util/_pywrap_determinism.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 enable(arg0: bool) -> None: ...
+def is_enabled() -> bool: ...

tensorflow/python/util/_pywrap_kernel_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 TryFindKernelClass(arg0: str) -> bytes: ...

tensorflow/python/util/_pywrap_nest.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 FlattenDictItems(arg0: object) -> object: ...

tensorflow/python/util/_pywrap_stat_summarizer.pyi

@@ -0,0 +1,26 @@
+# 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 overload
+
+class StatSummarizer:
+    @overload
+    def __init__(self, arg0: str) -> None: ...
+    @overload
+    def __init__(self) -> None: ...
+    def GetOutputString(self) -> str: ...
+    def PrintStepStats(self) -> None: ...
+    def ProcessStepStats(self, arg0) -> None: ...
+    def ProcessStepStatsStr(self, arg0: str) -> None: ...

tensorflow/python/util/_pywrap_tensor_float_32_execution.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 enable(arg0: bool) -> None: ...
+def is_enabled() -> bool: ...

tensorflow/python/util/_pywrap_tfprof.pyi

@@ -0,0 +1,23 @@
+# 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 AddStep(arg0: int, arg1: str, arg2: str, arg3: str) -> float: ...
+def DeleteProfiler() -> None: ...
+def NewProfiler(arg0: str, arg1: str) -> bool: ...
+def PrintModelAnalysis(arg0: str, arg1: str, arg2: str, arg3: str, arg4: str) -> bytes: ...
+def Profile(arg0: str, arg1: str) -> bytes: ...
+def ProfilerFromFile(arg0: str) -> None: ...
+def SerializeToString() -> bytes: ...
+def WriteProfile(arg0: str) -> None: ...

tensorflow/python/util/_pywrap_transform_graph.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 TransformGraphWithStringInputs(arg0: object, arg1: object, arg2: object, arg3: object) -> bytes: ...

tensorflow/python/util/_pywrap_util_port.pyi

@@ -0,0 +1,26 @@
+# 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 GpuSupportsHalfMatMulAndConv() -> bool: ...
+def IsAArch32Available() -> bool: ...
+def IsAArch64Available() -> bool: ...
+def IsBuiltWithNvcc() -> bool: ...
+def IsBuiltWithROCm() -> bool: ...
+def IsBuiltWithXLA() -> bool: ...
+def IsGoogleCudaEnabled() -> bool: ...
+def IsMklEnabled() -> bool: ...
+def IsPowerPCAvailable() -> bool: ...
+def IsSystemZAvailable() -> bool: ...
+def IsX86Available() -> bool: ...

tensorflow/python/util/_pywrap_utils.pyi

@@ -0,0 +1,35 @@
+# 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 AssertSameStructure(arg0: object, arg1: object, arg2: bool, arg3: bool) -> bool: ...
+def AssertSameStructureForData(arg0: object, arg1: object, arg2: bool) -> bool: ...
+def Flatten(arg0: object, arg1: bool) -> object: ...
+def FlattenForData(arg0: object) -> object: ...
+def IsAttrs(arg0: object) -> bool: ...
+def IsCompositeTensor(arg0: object) -> bool: ...
+def IsDataTypeSupportedByOneDNNOnThisCPU(arg0) -> bool: ...
+def IsMapping(arg0: object) -> bool: ...
+def IsMappingView(arg0: object) -> bool: ...
+def IsMutableMapping(arg0: object) -> bool: ...
+def IsNamedtuple(arg0: object, arg1: bool) -> object: ...
+def IsNested(arg0: object) -> bool: ...
+def IsNestedForData(arg0: object) -> bool: ...
+def IsNestedOrComposite(arg0: object) -> bool: ...
+def IsResourceVariable(arg0: object) -> bool: ...
+def IsTensor(arg0: object) -> bool: ...
+def IsTypeSpec(arg0: object) -> bool: ...
+def IsVariable(arg0: object) -> bool: ...
+def RegisterPyObject(arg0: object, arg1: object) -> object: ...
+def SameNamedtuples(arg0: object, arg1: object) -> object: ...

tensorflow/python/util/_tf_stack.pyi

@@ -0,0 +1,64 @@
+# 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.
+# ==============================================================================
+
+import typing
+from typing import Iterator, overload
+
+class GraphDebugInfoBuilder:
+    def __init__(self) -> None: ...
+    def AccumulateStackTrace(self, function: str, op: str, trace) -> None: ...
+    def AppendGraphDebugInfo(self, prefix: str, debug_info: bytes) -> None: ...
+    def Build(self) -> bytes: ...
+
+class PyBindFileSet:
+    def __init__(self) -> None: ...
+    def update_to(self, arg0: set) -> None: ...
+
+class PyBindSourceMap:
+    def __init__(self) -> None: ...
+    def update_to(self, arg0: tuple) -> None: ...
+
+class StackFrame:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def __eq__(self, arg0: StackFrame) -> bool: ...
+    def __getitem__(self, arg0: object) -> object: ...
+    def __hash__(self) -> int: ...
+    def __iter__(self) -> Iterator: ...
+    def __len__(self) -> int: ...
+    def __ne__(self, arg0: StackFrame) -> bool: ...
+    @property
+    def filename(self) -> str: ...
+    @property
+    def line(self) -> str: ...
+    @property
+    def lineno(self) -> int: ...
+    @property
+    def name(self) -> str: ...
+
+class StackTrace:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def get_user_frames(self) -> StackTrace: ...
+    def last_user_frame(self) -> StackFrame: ...
+    def __eq__(self, arg0: StackTrace) -> bool: ...
+    @overload
+    def __getitem__(self, arg0: int) -> StackFrame: ...
+    @overload
+    def __getitem__(self, arg0: slice) -> StackTrace: ...
+    def __iter__(self) -> typing.Iterator[StackFrame]: ...
+    def __hash__(self) -> int: ...
+    def __len__(self) -> int: ...
+
+def LoadTracesFromDebugInfo(debug_info_proto: bytes) -> dict[str, StackTrace]: ...
+def extract_stack(source_map: PyBindSourceMap, file_set: PyBindFileSet, stacklevel: int = ...) -> StackTrace: ...

tensorflow/python/util/all_util.py

@@ -0,0 +1,119 @@
+# 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.
+# ==============================================================================
+
+"""Generate __all__ from a module docstring."""
+import re as _re
+import sys as _sys
+
+from tensorflow.python.util import tf_inspect as _tf_inspect
+
+
+_reference_pattern = _re.compile(r'^@@(\w+)$', flags=_re.MULTILINE)
+
+
+def make_all(module_name, doc_string_modules=None):
+  """Generates `__all__` from the docstring of one or more modules.
+
+  Usage: `make_all(__name__)` or
+  `make_all(__name__, [sys.modules(__name__), other_module])`. The doc string
+  modules must each a docstring, and `__all__` will contain all symbols with
+  `@@` references, where that symbol currently exists in the module named
+  `module_name`.
+
+  Args:
+    module_name: The name of the module (usually `__name__`).
+    doc_string_modules: a list of modules from which to take docstring.
+    If None, then a list containing only the module named `module_name` is used.
+
+  Returns:
+    A list suitable for use as `__all__`.
+  """
+  if doc_string_modules is None:
+    doc_string_modules = [_sys.modules[module_name]]
+  cur_members = set(
+      name for name, _ in _tf_inspect.getmembers(_sys.modules[module_name]))
+
+  results = set()
+  for doc_module in doc_string_modules:
+    if doc_module.__doc__ is None:
+      continue
+    results.update([m.group(1)
+                    for m in _reference_pattern.finditer(doc_module.__doc__)
+                    if m.group(1) in cur_members])
+  return list(results)
+
+# Hidden attributes are attributes that have been hidden by
+# `remove_undocumented`. They can be re-instated by `reveal_undocumented`.
+# This maps symbol names to a tuple, containing:
+#   (module object, attribute value)
+_HIDDEN_ATTRIBUTES = {}
+
+
+def reveal_undocumented(symbol_name, target_module=None):
+  """Reveals a symbol that was previously removed by `remove_undocumented`.
+
+  This should be used by tensorflow internal tests only. It explicitly
+  defeats the encapsulation afforded by `remove_undocumented`.
+
+  It throws an exception when the symbol was not hidden in the first place.
+
+  Args:
+    symbol_name: a string representing the full absolute path of the symbol.
+    target_module: if specified, the module in which to restore the symbol.
+  """
+  if symbol_name not in _HIDDEN_ATTRIBUTES:
+    raise LookupError('Symbol %s is not a hidden symbol' % symbol_name)
+  symbol_basename = symbol_name.split('.')[-1]
+  (original_module, attr_value) = _HIDDEN_ATTRIBUTES[symbol_name]
+  if not target_module: target_module = original_module
+  setattr(target_module, symbol_basename, attr_value)
+
+
+def remove_undocumented(module_name, allowed_exception_list=None,
+                        doc_string_modules=None):
+  """Removes symbols in a module that are not referenced by a docstring.
+
+  Args:
+    module_name: the name of the module (usually `__name__`).
+    allowed_exception_list: a list of names that should not be removed.
+    doc_string_modules: a list of modules from which to take the docstrings.
+    If None, then a list containing only the module named `module_name` is used.
+
+    Furthermore, if a symbol previously added with `add_to_global_allowlist`,
+    then it will always be allowed. This is useful for internal tests.
+
+  Returns:
+    None
+  """
+  current_symbols = set(dir(_sys.modules[module_name]))
+  should_have = make_all(module_name, doc_string_modules)
+  should_have += allowed_exception_list or []
+  extra_symbols = current_symbols - set(should_have)
+  target_module = _sys.modules[module_name]
+  for extra_symbol in extra_symbols:
+    # Skip over __file__, etc. Also preserves internal symbols.
+    if extra_symbol.startswith('_'): continue
+    fully_qualified_name = module_name + '.' + extra_symbol
+    _HIDDEN_ATTRIBUTES[fully_qualified_name] = (target_module,
+                                                getattr(target_module,
+                                                        extra_symbol))
+    delattr(target_module, extra_symbol)
+
+
+__all__ = [
+    'make_all',
+    'remove_undocumented',
+    'reveal_undocumented',
+]

tensorflow/python/util/compat.py

@@ -0,0 +1,230 @@
+# 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.
+# ==============================================================================
+"""Compatibility functions.
+
+The `tf.compat` module contains two sets of compatibility functions.
+
+## Tensorflow 1.x and 2.x APIs
+
+The `compat.v1` and `compat.v2` submodules provide a complete copy of both the
+`v1` and `v2` APIs for backwards and forwards compatibility across TensorFlow
+versions 1.x and 2.x. See the
+[migration guide](https://www.tensorflow.org/guide/migrate) for details.
+
+## Utilities for writing compatible code
+
+Aside from the `compat.v1` and `compat.v2` submodules, `tf.compat` also contains
+a set of helper functions for writing code that works in both:
+
+* TensorFlow 1.x and 2.x
+* Python 2 and 3
+
+
+## Type collections
+
+The compatibility module also provides the following aliases for common
+sets of python types:
+
+* `bytes_or_text_types`
+* `complex_types`
+* `integral_types`
+* `real_types`
+
+API docstring: tensorflow.compat
+"""
+
+import codecs
+import collections.abc as collections_abc  # pylint: disable=unused-import
+import numbers as _numbers
+
+import numpy as _np
+
+from tensorflow.python.util.tf_export import tf_export
+
+
+def as_bytes(bytes_or_text, encoding='utf-8'):
+  """Converts `bytearray`, `bytes`, or unicode python input types to `bytes`.
+
+  Uses utf-8 encoding for text by default.
+
+  Args:
+    bytes_or_text: A `bytearray`, `bytes`, `str`, or `unicode` object.
+    encoding: A string indicating the charset for encoding unicode.
+
+  Returns:
+    A `bytes` object.
+
+  Raises:
+    TypeError: If `bytes_or_text` is not a binary or unicode string.
+  """
+  # Validate encoding, a LookupError will be raised if invalid.
+  encoding = codecs.lookup(encoding).name
+  if isinstance(bytes_or_text, bytearray):
+    return bytes(bytes_or_text)
+  elif isinstance(bytes_or_text, str):
+    return bytes_or_text.encode(encoding)
+  elif isinstance(bytes_or_text, bytes):
+    return bytes_or_text
+  else:
+    try:
+      return bytes(memoryview(bytes_or_text))
+    except TypeError as e:
+      raise TypeError(
+          'Expected binary or unicode string, got %r' % (bytes_or_text,)
+      ) from e
+
+
+def as_text(bytes_or_text, encoding='utf-8'):
+  """Converts any string-like python input types to unicode.
+
+  Returns the input as a unicode string. Uses utf-8 encoding for text
+  by default.
+
+  Args:
+    bytes_or_text: A `bytes`, `str`, or `unicode` object.
+    encoding: A string indicating the charset for decoding unicode.
+
+  Returns:
+    A `unicode` (Python 2) or `str` (Python 3) object.
+
+  Raises:
+    TypeError: If `bytes_or_text` is not a binary or unicode string.
+  """
+  # Validate encoding, a LookupError will be raised if invalid.
+  encoding = codecs.lookup(encoding).name
+  if isinstance(bytes_or_text, str):
+    return bytes_or_text
+  elif isinstance(bytes_or_text, bytes):
+    return bytes_or_text.decode(encoding)
+  else:
+    raise TypeError('Expected binary or unicode string, got %r' % bytes_or_text)
+
+
+def as_str(bytes_or_text, encoding='utf-8'):
+  """Acts as an alias for the `as_text` function..
+
+  Args:
+    bytes_or_text: The input value to be converted. A bytes or unicode object.
+    encoding: Optional string. The encoding to use if bytes_or_text is a bytes
+      object. Defaults to 'utf-8'.
+
+  Returns:
+    A unicode string.
+
+  Raises:
+    TypeError: If bytes_or_text is not a bytes or unicode object.
+    UnicodeDecodeError: If bytes_or_text is a bytes object and cannot be
+                        decoded using the specified encoding.
+  """
+  return as_text(bytes_or_text, encoding)
+
+tf_export('compat.as_text')(as_text)
+tf_export('compat.as_bytes')(as_bytes)
+tf_export('compat.as_str')(as_str)
+
+
+@tf_export('compat.as_str_any')
+def as_str_any(value, encoding='utf-8'):
+  """Converts input to `str` type.
+
+     Uses `str(value)`, except for `bytes` typed inputs, which are converted
+     using `as_str`.
+
+  Args:
+    value: A object that can be converted to `str`.
+    encoding: Encoding for `bytes` typed inputs.
+
+  Returns:
+    A `str` object.
+  """
+  if isinstance(value, bytes):
+    return as_str(value, encoding=encoding)
+  else:
+    return str(value)
+
+
+@tf_export('compat.path_to_str')
+def path_to_str(path):
+  r"""Converts input which is a `PathLike` object to `str` type.
+
+  Converts from any python constant representation of a `PathLike` object to
+  a string. If the input is not a `PathLike` object, simply returns the input.
+
+  Args:
+    path: An object that can be converted to path representation.
+
+  Returns:
+    A `str` object.
+
+  Usage:
+    In case a simplified `str` version of the path is needed from an
+    `os.PathLike` object.
+
+  Examples:
+  ```python
+  $ tf.compat.path_to_str('C:\XYZ\tensorflow\./.././tensorflow')
+  'C:\XYZ\tensorflow\./.././tensorflow' # Windows OS
+  $ tf.compat.path_to_str(Path('C:\XYZ\tensorflow\./.././tensorflow'))
+  'C:\XYZ\tensorflow\..\tensorflow' # Windows OS
+  $ tf.compat.path_to_str(Path('./corpus'))
+  'corpus' # Linux OS
+  $ tf.compat.path_to_str('./.././Corpus')
+  './.././Corpus' # Linux OS
+  $ tf.compat.path_to_str(Path('./.././Corpus'))
+  '../Corpus' # Linux OS
+  $ tf.compat.path_to_str(Path('./..////../'))
+  '../..' # Linux OS
+
+  ```
+  """
+  if hasattr(path, '__fspath__'):
+    path = as_str_any(path.__fspath__())
+  return path
+
+
+def path_to_bytes(path):
+  r"""Converts input which is a `PathLike` object to `bytes`.
+
+  Converts from any python constant representation of a `PathLike` object
+  or `str` to bytes.
+
+  Args:
+    path: An object that can be converted to path representation.
+
+  Returns:
+    A `bytes` object.
+
+  Usage:
+    In case a simplified `bytes` version of the path is needed from an
+    `os.PathLike` object.
+  """
+  if hasattr(path, '__fspath__'):
+    path = path.__fspath__()
+  return as_bytes(path)
+
+
+# Numpy 1.8 scalars don't inherit from numbers.Integral in Python 3, so we
+# need to check them specifically.  The same goes from Real and Complex.
+integral_types = (_numbers.Integral, _np.integer)
+tf_export('compat.integral_types').export_constant(__name__, 'integral_types')
+real_types = (_numbers.Real, _np.integer, _np.floating)
+tf_export('compat.real_types').export_constant(__name__, 'real_types')
+complex_types = (_numbers.Complex, _np.number)
+tf_export('compat.complex_types').export_constant(__name__, 'complex_types')
+
+# Either bytes or text.
+bytes_or_text_types = (bytes, str)
+tf_export('compat.bytes_or_text_types').export_constant(__name__,
+                                                        'bytes_or_text_types')

tensorflow/python/util/compat_test.py

@@ -0,0 +1,50 @@
+# 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.
+# ==============================================================================
+"""Compat tests."""
+
+import pickle
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class CompatTest(test.TestCase):
+
+  def testCompatValidEncoding(self):
+    self.assertEqual(compat.as_bytes("hello", "utf8"), b"hello")
+    self.assertEqual(compat.as_text(b"hello", "utf-8"), "hello")
+
+  def testCompatInvalidEncoding(self):
+    with self.assertRaises(LookupError):
+      compat.as_bytes("hello", "invalid")
+
+    with self.assertRaises(LookupError):
+      compat.as_text(b"hello", "invalid")
+
+  def testCompatBufferObjects(self):
+    pb = pickle.PickleBuffer(b"hello")
+    self.assertEqual(compat.as_bytes(pb), b"hello")
+    mv = memoryview(b"world")
+    self.assertEqual(compat.as_bytes(mv), b"world")
+
+  def testCompatInvalidObjects(self):
+    with self.assertRaises(TypeError):
+      compat.as_bytes(5)
+    with self.assertRaises(TypeError):
+      compat.as_bytes([1, 2, 3])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/custom_nest_protocol.py

@@ -0,0 +1,120 @@
+# 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.
+# ==============================================================================
+
+"""Protocol class for custom tf.nest support."""
+
+import typing
+from typing import Protocol
+
+
+@typing.runtime_checkable
+class CustomNestProtocol(Protocol):
+  """Protocol for adding custom tf.nest support in user-defined classes.
+
+  User classes should implement the two methods defined in this protocol in
+  order to be supported by nest functions.
+    - `__tf_flatten__` for generating the flattened components and the metadata
+      of the current object.
+    - `__tf_unflatten__` for creating a new object based on the input metadata
+      and the components.
+  See the method doc for details.
+
+  In terms of support level, classes implementing this protocol
+    - are supported by tf.nest and tf.data functions.
+    - have limited support from tf.function, which requires writing a custom
+      TraceType subclass to be used as the input or output of a tf.function.
+    - are NOT supported by SavedModel.
+
+  Code Examples:
+
+  >>> import dataclasses
+  >>> @dataclasses.dataclass
+  ... class MaskedTensor:
+  ...   mask: bool
+  ...   value: tf.Tensor
+  ...
+  ...   def __tf_flatten__(self):
+  ...     metadata = (self.mask,)  # static config.
+  ...     components = (self.value,)  # dynamic values.
+  ...     return metadata, components
+  ...
+  ...   @classmethod
+  ...   def __tf_unflatten__(cls, metadata, components):
+  ...     mask = metadata[0]
+  ...     value = components[0]
+  ...     return MaskedTensor(mask=mask, value=value)
+  ...
+  >>> mt = MaskedTensor(mask=True, value=tf.constant([1]))
+  >>> mt
+  MaskedTensor(mask=True, value=<tf.Tensor: ... numpy=array([1], dtype=int32)>)
+  >>> tf.nest.is_nested(mt)
+  True
+  >>> mt2 = MaskedTensor(mask=False, value=tf.constant([2]))
+  >>> tf.nest.assert_same_structure(mt, mt2)
+
+  >>> leaves = tf.nest.flatten(mt)
+  >>> leaves
+  [<tf.Tensor: shape=(1,), dtype=int32, numpy=array([1], dtype=int32)>]
+
+  >>> mt3 = tf.nest.pack_sequence_as(mt, leaves)
+  >>> mt3
+  MaskedTensor(mask=True, value=<tf.Tensor: ... numpy=array([1], dtype=int32)>)
+  >>> bool(mt == mt3)
+  True
+
+  >>> tf.nest.map_structure(lambda x: x * 2, mt)
+  MaskedTensor(mask=True, value=<tf.Tensor: ... numpy=array([2], dtype=int32)>)
+
+  More examples are available in the unit tests (nest_test.py).
+  """
+
+  def __tf_flatten__(self):
+    """Flatten current object into (metadata, components).
+
+    Returns:
+      A `tuple` of (metadata, components), where
+        - metadata is a custom Python object that stands for the static config
+          of the current object, which is supposed to be fixed and not affected
+          by data transformation.
+        - components is a `tuple` that contains the modifiable fields of the
+          current object.
+
+    Implementation Note:
+    - This method should not invoke any TensorFlow ops.
+    - This method only needs to flatten the current level. If current object has
+      an attribute that also need custom flattening, nest functions (such as
+      `nest.flatten`) will utilize this method to do recursive flattening.
+    - Components must be a `tuple`, not a `list`
+    """
+
+  @classmethod
+  def __tf_unflatten__(cls, metadata, components):
+    """Create a user-defined object from (metadata, components).
+
+    Args:
+      metadata: a custom Python object that stands for the static config for
+        reconstructing a new object of the current class.
+      components: a `tuple` that contains the dynamic data fields of the current
+        class, for object reconstruction.
+
+    Returns:
+      The user-defined object, with the same class of the current object.
+
+    Implementation Note:
+    - This method should not invoke any TensorFlow ops.
+    - This method only needs to unflatten the current level. If the object has
+      an attribute that also need custom unflattening, nest functions will
+      utilize this method to do recursive unflattening.
+    """

tensorflow/python/util/decorator_utils.py

@@ -0,0 +1,203 @@
+# 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.
+# ==============================================================================
+
+"""Utility functions for writing decorators (which modify docstrings)."""
+import sys
+
+
+def get_qualified_name(function):
+  # Python 3
+  if hasattr(function, '__qualname__'):
+    return function.__qualname__
+
+  # Python 2
+  if hasattr(function, 'im_class'):
+    return function.im_class.__name__ + '.' + function.__name__
+  return function.__name__
+
+
+def _normalize_docstring(docstring):
+  """Normalizes the docstring.
+
+  Replaces tabs with spaces, removes leading and trailing blanks lines, and
+  removes any indentation.
+
+  Copied from PEP-257:
+  https://www.python.org/dev/peps/pep-0257/#handling-docstring-indentation
+
+  Args:
+    docstring: the docstring to normalize
+
+  Returns:
+    The normalized docstring
+  """
+  if not docstring:
+    return ''
+  # Convert tabs to spaces (following the normal Python rules)
+  # and split into a list of lines:
+  lines = docstring.expandtabs().splitlines()
+  # Determine minimum indentation (first line doesn't count):
+  # (we use sys.maxsize because sys.maxint doesn't exist in Python 3)
+  indent = sys.maxsize
+  for line in lines[1:]:
+    stripped = line.lstrip()
+    if stripped:
+      indent = min(indent, len(line) - len(stripped))
+  # Remove indentation (first line is special):
+  trimmed = [lines[0].strip()]
+  if indent < sys.maxsize:
+    for line in lines[1:]:
+      trimmed.append(line[indent:].rstrip())
+  # Strip off trailing and leading blank lines:
+  while trimmed and not trimmed[-1]:
+    trimmed.pop()
+  while trimmed and not trimmed[0]:
+    trimmed.pop(0)
+  # Return a single string:
+  return '\n'.join(trimmed)
+
+
+def add_notice_to_docstring(doc,
+                            instructions,
+                            no_doc_str,
+                            suffix_str,
+                            notice,
+                            notice_type='Warning'):
+  """Adds a deprecation notice to a docstring.
+
+  Args:
+    doc: The original docstring.
+    instructions: A string, describing how to fix the problem.
+    no_doc_str: The default value to use for `doc` if `doc` is empty.
+    suffix_str: Is added to the end of the first line.
+    notice: A list of strings. The main notice warning body.
+    notice_type: The type of notice to use. Should be one of `[Caution,
+    Deprecated, Important, Note, Warning]`
+
+  Returns:
+    A new docstring, with the notice attached.
+
+  Raises:
+    ValueError: If `notice` is empty.
+  """
+  allowed_notice_types = ['Deprecated', 'Warning', 'Caution', 'Important',
+                          'Note']
+  if notice_type not in allowed_notice_types:
+    raise ValueError(
+        f'Unrecognized notice type. Should be one of: {allowed_notice_types}')
+
+  if not doc:
+    lines = [no_doc_str]
+  else:
+    lines = _normalize_docstring(doc).splitlines()
+    lines[0] += ' ' + suffix_str
+
+  if not notice:
+    raise ValueError('The `notice` arg must not be empty.')
+
+  notice[0] = f'{notice_type}: {notice[0]}'
+  notice = [''] + notice + ([instructions] if instructions else [])
+
+  if len(lines) > 1:
+    # Make sure that we keep our distance from the main body
+    if lines[1].strip():
+      notice.append('')
+
+    lines[1:1] = notice
+  else:
+    lines += notice
+
+  return '\n'.join(lines)
+
+
+def validate_callable(func, decorator_name):
+  if not hasattr(func, '__call__'):
+    raise ValueError(
+        '%s is not a function. If this is a property, make sure'
+        ' @property appears before @%s in your source code:'
+        '\n\n@property\n@%s\ndef method(...)' % (
+            func, decorator_name, decorator_name))
+
+
+class classproperty(object):  # pylint: disable=invalid-name
+  """Class property decorator.
+
+  Example usage:
+
+  class MyClass(object):
+
+    @classproperty
+    def value(cls):
+      return '123'
+
+  > print MyClass.value
+  123
+  """
+
+  def __init__(self, func):
+    self._func = func
+
+  def __get__(self, owner_self, owner_cls):
+    return self._func(owner_cls)
+
+
+class _CachedClassProperty(object):
+  """Cached class property decorator.
+
+  Transforms a class method into a property whose value is computed once
+  and then cached as a normal attribute for the life of the class.  Example
+  usage:
+
+  >>> class MyClass(object):
+  ...   @cached_classproperty
+  ...   def value(cls):
+  ...     print("Computing value")
+  ...     return '<property of %s>' % cls.__name__
+  >>> class MySubclass(MyClass):
+  ...   pass
+  >>> MyClass.value
+  Computing value
+  '<property of MyClass>'
+  >>> MyClass.value  # uses cached value
+  '<property of MyClass>'
+  >>> MySubclass.value
+  Computing value
+  '<property of MySubclass>'
+
+  This decorator is similar to `functools.cached_property`, but it adds a
+  property to the class, not to individual instances.
+  """
+
+  def __init__(self, func):
+    self._func = func
+    self._cache = {}
+
+  def __get__(self, obj, objtype):
+    if objtype not in self._cache:
+      self._cache[objtype] = self._func(objtype)
+    return self._cache[objtype]
+
+  def __set__(self, obj, value):
+    raise AttributeError('property %s is read-only' % self._func.__name__)
+
+  def __delete__(self, obj):
+    raise AttributeError('property %s is read-only' % self._func.__name__)
+
+
+def cached_classproperty(func):
+  return _CachedClassProperty(func)
+
+
+cached_classproperty.__doc__ = _CachedClassProperty.__doc__

tensorflow/python/util/decorator_utils_test.py

@@ -0,0 +1,166 @@
+# 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.
+# ==============================================================================
+"""decorator_utils tests."""
+
+# pylint: disable=unused-import
+import functools
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import decorator_utils
+
+
+def _test_function(unused_arg=0):
+  pass
+
+
+class GetQualifiedNameTest(test.TestCase):
+
+  def test_method(self):
+    self.assertEqual(
+        "GetQualifiedNameTest.test_method",
+        decorator_utils.get_qualified_name(GetQualifiedNameTest.test_method))
+
+  def test_function(self):
+    self.assertEqual("_test_function",
+                     decorator_utils.get_qualified_name(_test_function))
+
+
+class AddNoticeToDocstringTest(test.TestCase):
+
+  def _check(self, doc, expected):
+    self.assertEqual(
+        decorator_utils.add_notice_to_docstring(
+            doc=doc,
+            instructions="Instructions",
+            no_doc_str="Nothing here",
+            suffix_str="(suffix)",
+            notice=["Go away"]),
+        expected)
+
+  def test_regular(self):
+    expected = (
+        "Brief (suffix)\n\nWarning: Go away\nInstructions\n\nDocstring\n\n"
+        "Args:\n  arg1: desc")
+    # No indent for main docstring
+    self._check("Brief\n\nDocstring\n\nArgs:\n  arg1: desc", expected)
+    # 2 space indent for main docstring, blank lines not indented
+    self._check("Brief\n\n  Docstring\n\n  Args:\n    arg1: desc", expected)
+    # 2 space indent for main docstring, blank lines indented as well.
+    self._check("Brief\n  \n  Docstring\n  \n  Args:\n    arg1: desc", expected)
+    # No indent for main docstring, first line blank.
+    self._check("\n  Brief\n  \n  Docstring\n  \n  Args:\n    arg1: desc",
+                expected)
+    # 2 space indent, first line blank.
+    self._check("\n  Brief\n  \n  Docstring\n  \n  Args:\n    arg1: desc",
+                expected)
+
+  def test_brief_only(self):
+    expected = "Brief (suffix)\n\nWarning: Go away\nInstructions"
+    self._check("Brief", expected)
+    self._check("Brief\n", expected)
+    self._check("Brief\n  ", expected)
+    self._check("\nBrief\n  ", expected)
+    self._check("\n  Brief\n  ", expected)
+
+  def test_no_docstring(self):
+    expected = "Nothing here\n\nWarning: Go away\nInstructions"
+    self._check(None, expected)
+    self._check("", expected)
+
+  def test_no_empty_line(self):
+    expected = "Brief (suffix)\n\nWarning: Go away\nInstructions\n\nDocstring"
+    # No second line indent
+    self._check("Brief\nDocstring", expected)
+    # 2 space second line indent
+    self._check("Brief\n  Docstring", expected)
+    # No second line indent, first line blank
+    self._check("\nBrief\nDocstring", expected)
+    # 2 space second line indent, first line blank
+    self._check("\n  Brief\n  Docstring", expected)
+
+
+class ValidateCallableTest(test.TestCase):
+
+  def test_function(self):
+    decorator_utils.validate_callable(_test_function, "test")
+
+  def test_method(self):
+    decorator_utils.validate_callable(self.test_method, "test")
+
+  def test_callable(self):
+
+    class TestClass(object):
+
+      def __call__(self):
+        pass
+
+    decorator_utils.validate_callable(TestClass(), "test")
+
+  def test_partial(self):
+    partial = functools.partial(_test_function, unused_arg=7)
+    decorator_utils.validate_callable(partial, "test")
+
+  def test_fail_non_callable(self):
+    x = 0
+    self.assertRaises(ValueError, decorator_utils.validate_callable, x, "test")
+
+
+class CachedClassPropertyTest(test.TestCase):
+
+  def testCachedClassProperty(self):
+    log = []  # log all calls to `MyClass.value`.
+
+    class MyClass(object):
+
+      @decorator_utils.cached_classproperty
+      def value(cls):  # pylint: disable=no-self-argument
+        log.append(cls)
+        return cls.__name__
+
+    class MySubclass(MyClass):
+      pass
+
+    # Property is computed first time it is accessed.
+    self.assertLen(log, 0)
+    self.assertEqual(MyClass.value, "MyClass")
+    self.assertEqual(log, [MyClass])
+
+    # Cached values are used on subsequent accesses.
+    self.assertEqual(MyClass.value, "MyClass")
+    self.assertEqual(MyClass.value, "MyClass")
+    self.assertEqual(log, [MyClass])
+
+    # The wrapped method is called for each subclass.
+    self.assertEqual(MySubclass.value, "MySubclass")
+    self.assertEqual(log, [MyClass, MySubclass])
+    self.assertEqual(MySubclass.value, "MySubclass")
+    self.assertEqual(MySubclass.value, "MySubclass")
+    self.assertEqual(log, [MyClass, MySubclass])
+
+    # The property can also be accessed via an instance.
+    self.assertEqual(MyClass().value, "MyClass")
+    self.assertEqual(MySubclass().value, "MySubclass")
+    self.assertEqual(log, [MyClass, MySubclass])
+
+    # Attempts to modify the property via an instance will fail.
+    with self.assertRaises(AttributeError):
+      MyClass().value = 12
+    with self.assertRaises(AttributeError):
+      del MyClass().value
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/deprecated_module.py

@@ -0,0 +1,24 @@
+# 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.
+# ==============================================================================
+"""A deprecated module.
+
+For testing `deprecation.deprecate_moved_module`.
+"""
+
+from tensorflow.python.util import deprecated_module_new
+from tensorflow.python.util import deprecation
+
+__getattr__ = deprecation.deprecate_moved_module(
+    __name__, deprecated_module_new, "2.9")

tensorflow/python/util/deprecated_module_new.py

@@ -0,0 +1,22 @@
+# 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.
+# ==============================================================================
+"""A module to replace deprecated_module.
+
+For testing `deprecation.deprecate_moved_module`.
+"""
+
+
+def a():
+  return 1

tensorflow/python/util/deprecation.py

@@ -0,0 +1,765 @@
+# 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.
+# ==============================================================================
+"""Tensor utility functions."""
+import collections
+import functools
+import inspect
+import re
+
+from tensorflow.python.framework import strict_mode
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import decorator_utils
+from tensorflow.python.util import is_in_graph_mode
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+from tensorflow.tools.docs import doc_controls
+
+# Allow deprecation warnings to be silenced temporarily with a context manager.
+_PRINT_DEPRECATION_WARNINGS = True
+
+# Remember which deprecation warnings have been printed already.
+_PRINTED_WARNING = {}
+
+
+class DeprecatedNamesAlreadySetError(Exception):
+  """Raised when setting deprecated names multiple times for the same symbol."""
+
+
+def _log_deprecation(msg, *args, **kwargs):
+  """Raises errors for deprecated methods if in strict mode, warns otherwise."""
+  if strict_mode.STRICT_MODE:
+    logging.error(msg, *args, **kwargs)
+    raise RuntimeError(
+        'This behavior has been deprecated, which raises an error in strict'
+        ' mode.'
+    )
+  else:
+    logging.warning(msg, *args, **kwargs)
+
+
+def _add_deprecated_function_notice_to_docstring(doc, date, instructions):
+  """Adds a deprecation notice to a docstring for deprecated functions."""
+  main_text = [
+      'THIS FUNCTION IS DEPRECATED. It will be removed %s.'
+      % ('in a future version' if date is None else ('after %s' % date))
+  ]
+  if instructions:
+    main_text.append('Instructions for updating:')
+  return decorator_utils.add_notice_to_docstring(
+      doc,
+      instructions,
+      'DEPRECATED FUNCTION',
+      '(deprecated)',
+      main_text,
+      notice_type='Deprecated')
+
+
+def _add_deprecated_arg_notice_to_docstring(doc, date, instructions,
+                                            deprecated_names):
+  """Adds a deprecation notice to a docstring for deprecated arguments."""
+
+  deprecation_string = ', '.join(sorted(deprecated_names))
+
+  return decorator_utils.add_notice_to_docstring(
+      doc,
+      instructions,
+      'DEPRECATED FUNCTION ARGUMENTS',
+      '(deprecated arguments)', [
+          'SOME ARGUMENTS ARE DEPRECATED: `(%s)`. '
+          'They will be removed %s.' %
+          (deprecation_string, 'in a future version' if date is None else
+           ('after %s' % date)), 'Instructions for updating:'
+      ],
+      notice_type='Deprecated')
+
+
+def _add_deprecated_arg_value_notice_to_docstring(doc, date, instructions,
+                                                  deprecated_name_value_dict):
+  """Adds a deprecation notice to a docstring for deprecated arguments."""
+
+  deprecation_string = ', '.join(
+      '%s=%r' % (key, value)
+      for key, value in sorted(deprecated_name_value_dict.items()))
+
+  when = 'in a future version' if date is None else ('after %s' % date)
+
+  return decorator_utils.add_notice_to_docstring(
+      doc,
+      instructions,
+      'DEPRECATED FUNCTION ARGUMENT VALUES',
+      '(deprecated argument values)', [
+          'SOME ARGUMENT VALUES ARE DEPRECATED: `(%s)`. '
+          'They will be removed %s.' %
+          (deprecation_string, when), 'Instructions for updating:'
+      ],
+      notice_type='Deprecated')
+
+
+def _validate_deprecation_args(date, instructions):
+  if date is not None and not re.match(r'20\d\d-[01]\d-[0123]\d', date):
+    raise ValueError(f'Date must be in format YYYY-MM-DD. Received: {date}')
+  if not instructions:
+    raise ValueError(
+        'Don\'t deprecate things without conversion instructions! Specify '
+        'the `instructions` argument.')
+
+
+def _call_location(outer=False):
+  """Returns call location given level up from current call."""
+  # Two up: <_call_location>, <_call_location's caller>
+  # tf_inspect is not required here. Please ignore the lint warning by adding
+  # DISABLE_IMPORT_INSPECT_CHECK=TRUE to your cl description. Using it caused
+  # test timeouts (b/189384061).
+  f = inspect.currentframe().f_back.f_back
+  parent = f and f.f_back
+  if outer and parent is not None:
+    f = parent
+  return '{}:{}'.format(f.f_code.co_filename, f.f_lineno)
+
+
+def _safe_eq(a, b):
+  if a is None or b is None:
+    return a is None and b is None
+  return a == b
+
+
+def _wrap_decorator(wrapped_function, decorator_name):
+  """Indicate that one function wraps another.
+
+  This decorator wraps a function using `tf_decorator.make_decorator`
+  so that doc generation scripts can pick up original function
+  signature.
+  It would be better to use @functools.wrap decorator, but it would
+  not update function signature to match wrapped function in Python 2.
+
+  Args:
+    wrapped_function: The function that decorated function wraps.
+    decorator_name: The name of the decorator.
+
+  Returns:
+    Function that accepts wrapper function as an argument and returns
+    `TFDecorator` instance.
+  """
+
+  def wrapper(wrapper_func):
+    return tf_decorator.make_decorator(wrapped_function, wrapper_func,
+                                       decorator_name)
+
+  return wrapper
+
+
+def deprecated_alias(deprecated_name, name, func_or_class, warn_once=True):
+  """Deprecate a symbol in favor of a new name with identical semantics.
+
+  This function is meant to be used when defining a backwards-compatibility
+  alias for a symbol which has been moved. For example:
+
+  module1.py:
+  ```python
+  class NewNameForClass: pass
+  ```
+
+  module2.py:
+  ```python
+  import module1
+
+  DeprecatedNameForClass = deprecated_alias(
+    deprecated_name='module2.DeprecatedNameForClass',
+    name='module1.NewNameForClass',
+    func_or_class=module1.NewNameForClass)
+  ```
+
+  This function works for classes and functions.
+
+  For classes, it creates a new class which is functionally identical (it
+  inherits from the original, and overrides its constructor), but which prints
+  a deprecation warning when an instance is created. It also adds a deprecation
+  notice to the class' docstring.
+
+  For functions, it returns a function wrapped by `tf_decorator.make_decorator`.
+  That function prints a warning when used, and has a deprecation notice in its
+  docstring. This is more or less equivalent (the deprecation warning has
+  slightly different text) to writing:
+
+  ```python
+  @deprecated
+  def deprecated_alias(original_args):
+    real_function(original_args)
+  ```
+
+  Args:
+    deprecated_name: The name of the symbol that is being deprecated, to be used
+      in the warning message. This should be its fully qualified name to avoid
+      confusion.
+    name: The name of the symbol that is to be used instead of the deprecated
+      name. This should be a fully qualified name to avoid confusion.
+    func_or_class: The (non-deprecated) class or function for which a deprecated
+      alias should be created.
+    warn_once: If True (the default), only print a deprecation warning the first
+      time this function is used, or the class is instantiated.
+
+  Returns:
+    A wrapped version of `func_or_class` which prints a deprecation warning on
+    use and has a modified docstring.
+  """
+  if tf_inspect.isclass(func_or_class):
+
+    # Make a new class with __init__ wrapped in a warning.
+    class _NewClass(func_or_class):  # pylint: disable=missing-docstring
+      __doc__ = decorator_utils.add_notice_to_docstring(
+          func_or_class.__doc__,
+          'Please use %s instead.' % name,
+          'DEPRECATED CLASS',
+          '(deprecated)', [('THIS CLASS IS DEPRECATED. '
+                            'It will be removed in a future version. ')],
+          notice_type='Deprecated')
+      __name__ = func_or_class.__name__
+      __module__ = _call_location(outer=True)
+
+      @_wrap_decorator(func_or_class.__init__, 'deprecated_alias')
+      def __init__(self, *args, **kwargs):
+        if hasattr(_NewClass.__init__, '__func__'):
+          # Python 2
+          _NewClass.__init__.__func__.__doc__ = func_or_class.__init__.__doc__
+        else:
+          # Python 3
+          _NewClass.__init__.__doc__ = func_or_class.__init__.__doc__
+
+        if _PRINT_DEPRECATION_WARNINGS:
+          # We're making the alias as we speak. The original may have other
+          # aliases, so we cannot use it to check for whether it's already been
+          # warned about.
+          if _NewClass.__init__ not in _PRINTED_WARNING:
+            if warn_once:
+              _PRINTED_WARNING[_NewClass.__init__] = True
+            _log_deprecation(
+                'From %s: The name %s is deprecated. Please use %s instead.\n',
+                _call_location(), deprecated_name, name)
+        super(_NewClass, self).__init__(*args, **kwargs)
+
+    return _NewClass
+  else:
+    decorator_utils.validate_callable(func_or_class, 'deprecated')
+
+    # Make a wrapper for the original
+    @functools.wraps(func_or_class)
+    def new_func(*args, **kwargs):  # pylint: disable=missing-docstring
+      if _PRINT_DEPRECATION_WARNINGS:
+        # We're making the alias as we speak. The original may have other
+        # aliases, so we cannot use it to check for whether it's already been
+        # warned about.
+        if new_func not in _PRINTED_WARNING:
+          if warn_once:
+            _PRINTED_WARNING[new_func] = True
+          _log_deprecation(
+              'From %s: The name %s is deprecated. Please use %s instead.\n',
+              _call_location(), deprecated_name, name)
+      return func_or_class(*args, **kwargs)
+
+    return tf_decorator.make_decorator(
+        func_or_class, new_func, 'deprecated',
+        _add_deprecated_function_notice_to_docstring(
+            func_or_class.__doc__, None, 'Please use %s instead.' % name))
+
+
+def deprecated_endpoints(*args):
+  """Decorator for marking endpoints deprecated.
+
+  This decorator does not print deprecation messages.
+  TODO(annarev): eventually start printing deprecation warnings when
+  @deprecation_endpoints decorator is added.
+
+  Args:
+    *args: Deprecated endpoint names.
+
+  Returns:
+    A function that takes symbol as an argument and adds
+    _tf_deprecated_api_names to that symbol.
+    _tf_deprecated_api_names would be set to a list of deprecated
+    endpoint names for the symbol.
+  """
+
+  def deprecated_wrapper(func):
+    # pylint: disable=protected-access
+    if '_tf_deprecated_api_names' in func.__dict__:
+      raise DeprecatedNamesAlreadySetError(
+          f'Cannot set deprecated names for {func.__name__} to {args}. '
+          'Deprecated names are already set to '
+          f'{func._tf_deprecated_api_names}.')
+    func._tf_deprecated_api_names = args
+    # pylint: disable=protected-access
+    return func
+
+  return deprecated_wrapper
+
+
+def deprecated(date, instructions, warn_once=True):
+  """Decorator for marking functions or methods deprecated.
+
+  This decorator logs a deprecation warning whenever the decorated function is
+  called. It has the following format:
+
+    <function> (from <module>) is deprecated and will be removed after <date>.
+    Instructions for updating:
+    <instructions>
+
+  If `date` is None, 'after <date>' is replaced with 'in a future version'.
+  <function> will include the class name if it is a method.
+
+  It also edits the docstring of the function: ' (deprecated)' is appended
+  to the first line of the docstring and a deprecation notice is prepended
+  to the rest of the docstring.
+
+  Args:
+    date: String or None. The date the function is scheduled to be removed. Must
+      be ISO 8601 (YYYY-MM-DD), or None.
+    instructions: String. Instructions on how to update code using the
+      deprecated function.
+    warn_once: Boolean. Set to `True` to warn only the first time the decorated
+      function is called. Otherwise, every call will log a warning.
+
+  Returns:
+    Decorated function or method.
+
+  Raises:
+    ValueError: If date is not None or in ISO 8601 format, or instructions are
+      empty.
+  """
+  _validate_deprecation_args(date, instructions)
+
+  def deprecated_wrapper(func_or_class):
+    """Deprecation wrapper."""
+    if isinstance(func_or_class, type):
+      # If a class is deprecated, you actually want to wrap the constructor.
+      cls = func_or_class
+      if cls.__new__ is object.__new__:
+        # If a class defaults to its parent's constructor, wrap that instead.
+        func = cls.__init__
+        constructor_name = '__init__'
+        decorators, _ = tf_decorator.unwrap(func)
+        for decorator in decorators:
+          if decorator.decorator_name == 'deprecated':
+            # If the parent is already deprecated, there's nothing to do.
+            return cls
+      else:
+        func = cls.__new__
+        constructor_name = '__new__'
+
+    else:
+      cls = None
+      constructor_name = None
+      func = func_or_class
+
+    decorator_utils.validate_callable(func, 'deprecated')
+
+    @_wrap_decorator(func, 'deprecated')
+    def new_func(*args, **kwargs):  # pylint: disable=missing-docstring
+      if _PRINT_DEPRECATION_WARNINGS:
+        if func not in _PRINTED_WARNING and cls not in _PRINTED_WARNING:
+          if warn_once:
+            _PRINTED_WARNING[func] = True
+            if cls:
+              _PRINTED_WARNING[cls] = True
+          _log_deprecation(
+              'From %s: %s (from %s) is deprecated and will be removed %s.\n'
+              'Instructions for updating:\n%s', _call_location(),
+              decorator_utils.get_qualified_name(func),
+              func_or_class.__module__,
+              'in a future version' if date is None else ('after %s' % date),
+              instructions)
+      return func(*args, **kwargs)
+
+    doc_controls.set_deprecated(new_func)
+    new_func = tf_decorator.make_decorator(
+        func, new_func, 'deprecated',
+        _add_deprecated_function_notice_to_docstring(func.__doc__, date,
+                                                     instructions))
+    new_func.__signature__ = inspect.signature(func)
+
+    if cls is None:
+      return new_func
+    else:
+      # Insert the wrapped function as the constructor
+      setattr(cls, constructor_name, new_func)
+
+      # And update the docstring of the class.
+      cls.__doc__ = _add_deprecated_function_notice_to_docstring(
+          cls.__doc__, date, instructions)
+
+      return cls
+
+  return deprecated_wrapper
+
+
+DeprecatedArgSpec = collections.namedtuple(
+    'DeprecatedArgSpec', ['position', 'has_ok_value', 'ok_value'])
+
+
+def deprecated_args(date, instructions, *deprecated_arg_names_or_tuples,
+                    **kwargs):
+  """Decorator for marking specific function arguments as deprecated.
+
+  This decorator logs a deprecation warning whenever the decorated function is
+  called with the deprecated argument. It has the following format:
+
+    Calling <function> (from <module>) with <arg> is deprecated and will be
+    removed after <date>. Instructions for updating:
+      <instructions>
+
+  If `date` is None, 'after <date>' is replaced with 'in a future version'.
+  <function> includes the class name if it is a method.
+
+  It also edits the docstring of the function: ' (deprecated arguments)' is
+  appended to the first line of the docstring and a deprecation notice is
+  prepended to the rest of the docstring.
+
+  Args:
+    date: String or None. The date the function is scheduled to be removed. Must
+      be ISO 8601 (YYYY-MM-DD), or None.
+    instructions: String. Instructions on how to update code using the
+      deprecated function.
+    *deprecated_arg_names_or_tuples: String or 2-Tuple (String, ok_val).  The
+      string is the deprecated argument name. Optionally, an ok-value may be
+      provided.  If the user provided argument equals this value, the warning is
+      suppressed.
+    **kwargs: If `warn_once=False` is passed, every call with a deprecated
+      argument will log a warning. The default behavior is to only warn the
+      first time the function is called with any given deprecated argument. All
+      other kwargs raise `ValueError`.
+
+  Returns:
+    Decorated function or method.
+
+  Raises:
+    ValueError: If date is not None or in ISO 8601 format, instructions are
+      empty, the deprecated arguments are not present in the function
+      signature, the second element of a deprecated_tuple is not a
+      list, or if a kwarg other than `warn_once` is passed.
+  """
+  _validate_deprecation_args(date, instructions)
+  if not deprecated_arg_names_or_tuples:
+    raise ValueError('Specify which argument is deprecated.')
+  if kwargs and list(kwargs.keys()) != ['warn_once']:
+    kwargs.pop('warn_once', None)
+    raise ValueError(f'Illegal argument passed to deprecated_args: {kwargs}')
+  warn_once = kwargs.get('warn_once', True)
+
+  def _get_arg_names_to_ok_vals():
+    """Returns a dict mapping arg_name to DeprecatedArgSpec w/o position."""
+    d = {}
+    for name_or_tuple in deprecated_arg_names_or_tuples:
+      if isinstance(name_or_tuple, tuple):
+        d[name_or_tuple[0]] = DeprecatedArgSpec(-1, True, name_or_tuple[1])
+      else:
+        d[name_or_tuple] = DeprecatedArgSpec(-1, False, None)
+    return d
+
+  def _get_deprecated_positional_arguments(names_to_ok_vals, arg_spec):
+    """Builds a dictionary from deprecated arguments to their spec.
+
+    Returned dict is keyed by argument name.
+    Each value is a DeprecatedArgSpec with the following fields:
+       position: The zero-based argument position of the argument
+         within the signature.  None if the argument isn't found in
+         the signature.
+       ok_values:  Values of this argument for which warning will be
+         suppressed.
+
+    Args:
+      names_to_ok_vals: dict from string arg_name to a list of values, possibly
+        empty, which should not elicit a warning.
+      arg_spec: Output from tf_inspect.getfullargspec on the called function.
+
+    Returns:
+      Dictionary from arg_name to DeprecatedArgSpec.
+    """
+    # Extract argument list
+    arg_space = arg_spec.args + arg_spec.kwonlyargs
+    arg_name_to_pos = {name: pos for pos, name in enumerate(arg_space)}
+    deprecated_positional_args = {}
+    for arg_name, spec in iter(names_to_ok_vals.items()):
+      if arg_name in arg_name_to_pos:
+        pos = arg_name_to_pos[arg_name]
+        deprecated_positional_args[arg_name] = DeprecatedArgSpec(
+            pos, spec.has_ok_value, spec.ok_value)
+    return deprecated_positional_args
+
+  deprecated_arg_names = _get_arg_names_to_ok_vals()
+
+  def deprecated_wrapper(func):
+    """Deprecation decorator."""
+    decorator_utils.validate_callable(func, 'deprecated_args')
+
+    arg_spec = tf_inspect.getfullargspec(func)
+    deprecated_positions = _get_deprecated_positional_arguments(
+        deprecated_arg_names, arg_spec)
+
+    is_varargs_deprecated = arg_spec.varargs in deprecated_arg_names
+    is_kwargs_deprecated = arg_spec.varkw in deprecated_arg_names
+
+    if (len(deprecated_positions) + is_varargs_deprecated + is_kwargs_deprecated
+        != len(deprecated_arg_names_or_tuples)):
+      known_args = (
+          arg_spec.args + arg_spec.kwonlyargs +
+          [arg_spec.varargs, arg_spec.varkw])
+      missing_args = [
+          arg_name for arg_name in deprecated_arg_names
+          if arg_name not in known_args
+      ]
+      raise ValueError('The following deprecated arguments are not present '
+                       f'in the function signature: {missing_args}. '
+                       'Expected arguments from the following list: '
+                       f'{known_args}.')
+
+    def _same_value(a, b):
+      """A comparison operation that works for multiple object types.
+
+      Returns True for two empty lists, two numeric values with the
+      same value, etc.
+
+      Returns False for (pd.DataFrame, None), and other pairs which
+      should not be considered equivalent.
+
+      Args:
+        a: value one of the comparison.
+        b: value two of the comparison.
+
+      Returns:
+        A boolean indicating whether the two inputs are the same value
+        for the purposes of deprecation.
+      """
+      if a is b:
+        return True
+      try:
+        equality = a == b
+        if isinstance(equality, bool):
+          return equality
+      except TypeError:
+        return False
+      return False
+
+    @functools.wraps(func)
+    def new_func(*args, **kwargs):
+      """Deprecation wrapper."""
+      # TODO(apassos) figure out a way to have reasonable performance with
+      # deprecation warnings and eager mode.
+      if is_in_graph_mode.IS_IN_GRAPH_MODE() and _PRINT_DEPRECATION_WARNINGS:
+        invalid_args = []
+        named_args = tf_inspect.getcallargs(func, *args, **kwargs)
+        for arg_name, spec in iter(deprecated_positions.items()):
+          if (spec.position < len(args) and
+              not (spec.has_ok_value and
+                   _same_value(named_args[arg_name], spec.ok_value))):
+            invalid_args.append(arg_name)
+        if is_varargs_deprecated and len(args) > len(arg_spec.args):
+          invalid_args.append(arg_spec.varargs)
+        if is_kwargs_deprecated and kwargs:
+          invalid_args.append(arg_spec.varkw)
+        for arg_name in deprecated_arg_names:
+          if (arg_name in kwargs and
+              not (deprecated_positions[arg_name].has_ok_value and
+                   _same_value(named_args[arg_name],
+                               deprecated_positions[arg_name].ok_value))):
+            invalid_args.append(arg_name)
+        for arg_name in invalid_args:
+          if (func, arg_name) not in _PRINTED_WARNING:
+            if warn_once:
+              _PRINTED_WARNING[(func, arg_name)] = True
+            _log_deprecation(
+                'From %s: calling %s (from %s) with %s is deprecated and will '
+                'be removed %s.\nInstructions for updating:\n%s',
+                _call_location(), decorator_utils.get_qualified_name(func),
+                func.__module__, arg_name,
+                'in a future version' if date is None else ('after %s' % date),
+                instructions)
+      return func(*args, **kwargs)
+
+    doc = _add_deprecated_arg_notice_to_docstring(
+        func.__doc__, date, instructions, sorted(deprecated_arg_names.keys()))
+    return tf_decorator.make_decorator(func, new_func, 'deprecated', doc)
+
+  return deprecated_wrapper
+
+
+def deprecated_arg_values(date,
+                          instructions,
+                          warn_once=True,
+                          **deprecated_kwargs):
+  """Decorator for marking specific function argument values as deprecated.
+
+  This decorator logs a deprecation warning whenever the decorated function is
+  called with the deprecated argument values. It has the following format:
+
+    Calling <function> (from <module>) with <arg>=<value> is deprecated and
+    will be removed after <date>. Instructions for updating:
+      <instructions>
+
+  If `date` is None, 'after <date>' is replaced with 'in a future version'.
+  <function> will include the class name if it is a method.
+
+  It also edits the docstring of the function: ' (deprecated arguments)' is
+  appended to the first line of the docstring and a deprecation notice is
+  prepended to the rest of the docstring.
+
+  Args:
+    date: String or None. The date the function is scheduled to be removed. Must
+      be ISO 8601 (YYYY-MM-DD), or None
+    instructions: String. Instructions on how to update code using the
+      deprecated function.
+    warn_once: If `True`, warn only the first time this function is called with
+      deprecated argument values. Otherwise, every call (with a deprecated
+      argument value) will log a warning.
+    **deprecated_kwargs: The deprecated argument values.
+
+  Returns:
+    Decorated function or method.
+
+  Raises:
+    ValueError: If date is not None or in ISO 8601 format, or instructions are
+      empty.
+  """
+  _validate_deprecation_args(date, instructions)
+  if not deprecated_kwargs:
+    raise ValueError('Specify which argument values are deprecated.')
+
+  def deprecated_wrapper(func):
+    """Deprecation decorator."""
+    decorator_utils.validate_callable(func, 'deprecated_arg_values')
+
+    @functools.wraps(func)
+    def new_func(*args, **kwargs):
+      """Deprecation wrapper."""
+      if _PRINT_DEPRECATION_WARNINGS:
+        named_args = tf_inspect.getcallargs(func, *args, **kwargs)
+        for arg_name, arg_value in deprecated_kwargs.items():
+          if arg_name in named_args and _safe_eq(named_args[arg_name],
+                                                 arg_value):
+            if (func, arg_name) not in _PRINTED_WARNING:
+              if warn_once:
+                _PRINTED_WARNING[(func, arg_name)] = True
+              _log_deprecation(
+                  'From %s: calling %s (from %s) with %s=%s is deprecated and '
+                  'will be removed %s.\nInstructions for updating:\n%s',
+                  _call_location(), decorator_utils.get_qualified_name(func),
+                  func.__module__, arg_name, arg_value,
+                  'in a future version' if date is None else
+                  ('after %s' % date), instructions)
+      return func(*args, **kwargs)
+
+    doc = _add_deprecated_arg_value_notice_to_docstring(func.__doc__, date,
+                                                        instructions,
+                                                        deprecated_kwargs)
+    return tf_decorator.make_decorator(func, new_func, 'deprecated', doc)
+
+  return deprecated_wrapper
+
+
+def deprecated_argument_lookup(new_name, new_value, old_name, old_value):
+  """Looks up deprecated argument name and ensures both are not used.
+
+  Args:
+    new_name: new name of argument
+    new_value: value of new argument (or None if not used)
+    old_name: old name of argument
+    old_value: value of old argument (or None if not used)
+
+  Returns:
+    The effective argument that should be used.
+  Raises:
+    ValueError: if new_value and old_value are both non-null
+  """
+  if old_value is not None:
+    if new_value is not None:
+      raise ValueError(f"Cannot specify both '{old_name}' and '{new_name}'.")
+    return old_value
+  return new_value
+
+
+def rewrite_argument_docstring(old_doc, old_argument, new_argument):
+  if old_doc is None:
+    return None
+  return old_doc.replace('`%s`' % old_argument,
+                         '`%s`' % new_argument).replace('%s:' % old_argument,
+                                                        '%s:' % new_argument)
+
+
+@tf_contextlib.contextmanager
+def silence():
+  """Temporarily silence deprecation warnings."""
+  global _PRINT_DEPRECATION_WARNINGS
+  print_deprecation_warnings = _PRINT_DEPRECATION_WARNINGS
+  _PRINT_DEPRECATION_WARNINGS = False
+  yield
+  _PRINT_DEPRECATION_WARNINGS = print_deprecation_warnings
+
+
+def deprecate_moved_module(deprecated_name, new_module, deletion_version):
+  """Logs a warning when a module that has been moved to a new location is used.
+
+  Copy the following code into the old module:
+
+  ```
+  import deprecation
+  import new_module
+
+  __getattr__ = deprecation.deprecate_moved_module(
+      __name__, new_module, "2.9")  # adjust version number.
+  ```
+
+  Args:
+    deprecated_name: Name of old module.
+    new_module: Module to replace the old module.
+    deletion_version: Version of TensorFlow in which the old module will be
+      removed.
+
+  Returns:
+    A function that logs a warning and returns the symbol from the new module.
+    Set this function as the module's `__getattr__`.
+  """
+
+  def getter(name):
+    if getter not in _PRINTED_WARNING and _PRINT_DEPRECATION_WARNINGS:
+      _PRINTED_WARNING[getter] = True
+      _log_deprecation(
+          'Please fix your imports. Module %s has been moved to %s. The old '
+          'module will be deleted in version %s.', deprecated_name,
+          new_module.__name__, deletion_version)
+    return getattr(new_module, name)
+
+  return getter
+
+
+class HiddenTfApiAttribute(property):
+  """Hides a class attribute from the public API.
+
+  Attributes in public classes can be hidden from the API by having an '_' in
+  front of the name (e.g. ClassName._variables). This doesn't work when
+  attributes or methods are inherited from a parent class. To hide inherited
+  attributes, set their values to be `deprecation.hide_attribute_from_api`.
+  """
+
+  def __init__(self, deprecation_message):
+
+    def raise_error(unused_self):
+      raise AttributeError(deprecation_message)
+
+    super(HiddenTfApiAttribute, self).__init__(raise_error)
+
+
+hide_attribute_from_api = HiddenTfApiAttribute  # pylint: disable=invalid-name
+
+# TODO(kathywu): Remove once cl/246395236 is submitted.
+HIDDEN_ATTRIBUTE = HiddenTfApiAttribute('This attribute has been deprecated.')

tensorflow/python/util/determinism.cc

@@ -0,0 +1,23 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/util/determinism.h"
+
+#include "pybind11/pybind11.h"  // from @pybind11
+
+PYBIND11_MODULE(_pywrap_determinism, m) {
+  m.def("enable", &tensorflow::EnableOpDeterminism);
+  m.def("is_enabled", &tensorflow::OpDeterminismRequired);
+}

tensorflow/python/util/example_parser_configuration.py

@@ -0,0 +1,206 @@
+# 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.
+# ==============================================================================
+"""Extract parse_example op configuration to a proto."""
+
+from tensorflow.core.example import example_parser_configuration_pb2
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+
+
+def extract_example_parser_configuration(parse_example_op, sess):
+  """Returns an ExampleParserConfig proto.
+
+  Args:
+    parse_example_op: A ParseExample or ParseExampleV2 `Operation`
+    sess: A tf.compat.v1.Session needed to obtain some configuration values.
+  Returns:
+    A ExampleParserConfig proto.
+
+  Raises:
+    ValueError: If attributes are inconsistent.
+  """
+  if parse_example_op.type == "ParseExample":
+    return _extract_from_parse_example(parse_example_op, sess)
+  elif parse_example_op.type == "ParseExampleV2":
+    return _extract_from_parse_example_v2(parse_example_op, sess)
+  else:
+    raise ValueError(
+        "Found unexpected type when parsing example. Expected `ParseExample` "
+        f"object. Received type: {parse_example_op.type}")
+
+
+def _extract_from_parse_example(parse_example_op, sess):
+  """Extract ExampleParserConfig from ParseExample op."""
+  config = example_parser_configuration_pb2.ExampleParserConfiguration()
+
+  num_sparse = parse_example_op.get_attr("Nsparse")
+  num_dense = parse_example_op.get_attr("Ndense")
+  total_features = num_dense + num_sparse
+
+  sparse_types = parse_example_op.get_attr("sparse_types")
+  dense_types = parse_example_op.get_attr("Tdense")
+  dense_shapes = parse_example_op.get_attr("dense_shapes")
+
+  if len(sparse_types) != num_sparse:
+    raise ValueError("len(sparse_types) attribute does not match "
+                     "Nsparse attribute (%d vs %d)" %
+                     (len(sparse_types), num_sparse))
+
+  if len(dense_types) != num_dense:
+    raise ValueError("len(dense_types) attribute does not match "
+                     "Ndense attribute (%d vs %d)" %
+                     (len(dense_types), num_dense))
+
+  if len(dense_shapes) != num_dense:
+    raise ValueError("len(dense_shapes) attribute does not match "
+                     "Ndense attribute (%d vs %d)" %
+                     (len(dense_shapes), num_dense))
+
+  # Skip over the serialized input, and the names input.
+  fetch_list = parse_example_op.inputs[2:]
+
+  # Fetch total_features key names and num_dense default values.
+  if len(fetch_list) != (total_features + num_dense):
+    raise ValueError("len(fetch_list) does not match total features + "
+                     "num_dense (%d vs %d)" %
+                     (len(fetch_list), (total_features + num_dense)))
+
+  fetched = sess.run(fetch_list)
+
+  if len(fetched) != len(fetch_list):
+    raise ValueError("len(fetched) does not match len(fetch_list) "
+                     "(%d vs %d)" % (len(fetched), len(fetch_list)))
+
+  # Fetch indices.
+  sparse_keys_start = 0
+  dense_keys_start = sparse_keys_start + num_sparse
+  dense_def_start = dense_keys_start + num_dense
+
+  # Output tensor indices.
+  sparse_indices_start = 0
+  sparse_values_start = num_sparse
+  sparse_shapes_start = sparse_values_start + num_sparse
+  dense_values_start = sparse_shapes_start + num_sparse
+
+  # Dense features.
+  for i in range(num_dense):
+    key = fetched[dense_keys_start + i]
+    feature_config = config.feature_map[key]
+    # Convert the default value numpy array fetched from the session run
+    # into a TensorProto.
+    fixed_config = feature_config.fixed_len_feature
+
+    fixed_config.default_value.CopyFrom(
+        tensor_util.make_tensor_proto(fetched[dense_def_start + i]))
+    # Convert the shape from the attributes
+    # into a TensorShapeProto.
+    fixed_config.shape.CopyFrom(
+        tensor_shape.TensorShape(dense_shapes[i]).as_proto())
+
+    fixed_config.dtype = dense_types[i].as_datatype_enum
+    # Get the output tensor name.
+    fixed_config.values_output_tensor_name = parse_example_op.outputs[
+        dense_values_start + i].name
+
+  # Sparse features.
+  for i in range(num_sparse):
+    key = fetched[sparse_keys_start + i]
+    feature_config = config.feature_map[key]
+    var_len_feature = feature_config.var_len_feature
+    var_len_feature.dtype = sparse_types[i].as_datatype_enum
+    var_len_feature.indices_output_tensor_name = parse_example_op.outputs[
+        sparse_indices_start + i].name
+    var_len_feature.values_output_tensor_name = parse_example_op.outputs[
+        sparse_values_start + i].name
+    var_len_feature.shapes_output_tensor_name = parse_example_op.outputs[
+        sparse_shapes_start + i].name
+
+  return config
+
+
+def _extract_from_parse_example_v2(parse_example_op, sess):
+  """Extract ExampleParserConfig from ParseExampleV2 op."""
+  config = example_parser_configuration_pb2.ExampleParserConfiguration()
+
+  dense_types = parse_example_op.get_attr("Tdense")
+  num_sparse = parse_example_op.get_attr("num_sparse")
+  sparse_types = parse_example_op.get_attr("sparse_types")
+  ragged_value_types = parse_example_op.get_attr("ragged_value_types")
+  ragged_split_types = parse_example_op.get_attr("ragged_split_types")
+  dense_shapes = parse_example_op.get_attr("dense_shapes")
+
+  num_dense = len(dense_types)
+  num_ragged = len(ragged_value_types)
+  assert len(ragged_value_types) == len(ragged_split_types)
+  assert len(parse_example_op.inputs) == 5 + num_dense
+
+  # Skip over the serialized input, and the names input.
+  fetched = sess.run(parse_example_op.inputs[2:])
+  sparse_keys = fetched[0].tolist()
+  dense_keys = fetched[1].tolist()
+  ragged_keys = fetched[2].tolist()
+  dense_defaults = fetched[3:]
+  assert len(sparse_keys) == num_sparse
+  assert len(dense_keys) == num_dense
+  assert len(ragged_keys) == num_ragged
+
+  # Output tensor indices.
+  sparse_indices_start = 0
+  sparse_values_start = num_sparse
+  sparse_shapes_start = sparse_values_start + num_sparse
+  dense_values_start = sparse_shapes_start + num_sparse
+  ragged_values_start = dense_values_start + num_dense
+  ragged_row_splits_start = ragged_values_start + num_ragged
+
+  # Dense features.
+  for i in range(num_dense):
+    key = dense_keys[i]
+    feature_config = config.feature_map[key]
+    # Convert the default value numpy array fetched from the session run
+    # into a TensorProto.
+    fixed_config = feature_config.fixed_len_feature
+
+    fixed_config.default_value.CopyFrom(
+        tensor_util.make_tensor_proto(dense_defaults[i]))
+    # Convert the shape from the attributes
+    # into a TensorShapeProto.
+    fixed_config.shape.CopyFrom(
+        tensor_shape.TensorShape(dense_shapes[i]).as_proto())
+
+    fixed_config.dtype = dense_types[i].as_datatype_enum
+    # Get the output tensor name.
+    fixed_config.values_output_tensor_name = parse_example_op.outputs[
+        dense_values_start + i].name
+
+  # Sparse features.
+  for i in range(num_sparse):
+    key = sparse_keys[i]
+    feature_config = config.feature_map[key]
+    var_len_feature = feature_config.var_len_feature
+    var_len_feature.dtype = sparse_types[i].as_datatype_enum
+    var_len_feature.indices_output_tensor_name = parse_example_op.outputs[
+        sparse_indices_start + i].name
+    var_len_feature.values_output_tensor_name = parse_example_op.outputs[
+        sparse_values_start + i].name
+    var_len_feature.shapes_output_tensor_name = parse_example_op.outputs[
+        sparse_shapes_start + i].name
+
+  if num_ragged != 0:
+    del ragged_values_start  # unused
+    del ragged_row_splits_start  # unused
+    raise ValueError("Ragged features are not yet supported by "
+                     "example_parser_configuration.proto")
+
+  return config

tensorflow/python/util/example_parser_configuration_test.py

@@ -0,0 +1,94 @@
+# 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 ExampleParserConfiguration."""
+
+from google.protobuf import text_format
+
+from tensorflow.core.example import example_parser_configuration_pb2
+from tensorflow.python.client import session
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util.example_parser_configuration import extract_example_parser_configuration
+
+EXPECTED_CONFIG_V1 = """
+feature_map {
+  key: "x"
+  value {
+    fixed_len_feature {
+      dtype: DT_FLOAT
+      shape {
+        dim {
+          size: 1
+        }
+      }
+      default_value {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 1
+          }
+        }
+        float_val: 33.0
+      }
+      values_output_tensor_name: "ParseExample/ParseExample:3"
+    }
+  }
+}
+feature_map {
+  key: "y"
+  value {
+    var_len_feature {
+      dtype: DT_STRING
+      values_output_tensor_name: "ParseExample/ParseExample:1"
+      indices_output_tensor_name: "ParseExample/ParseExample:0"
+      shapes_output_tensor_name: "ParseExample/ParseExample:2"
+    }
+  }
+}
+"""
+
+
+EXPECTED_CONFIG_V2 = EXPECTED_CONFIG_V1.replace(
+    'ParseExample/ParseExample:', 'ParseExample/ParseExampleV2:')
+
+
+class ExampleParserConfigurationTest(test.TestCase):
+
+  def getExpectedConfig(self, op_type):
+    expected = example_parser_configuration_pb2.ExampleParserConfiguration()
+    if op_type == 'ParseExampleV2':
+      text_format.Parse(EXPECTED_CONFIG_V2, expected)
+    else:
+      text_format.Parse(EXPECTED_CONFIG_V1, expected)
+    return expected
+
+  def testBasic(self):
+    with session.Session() as sess:
+      examples = array_ops.placeholder(dtypes.string, shape=[1])
+      feature_to_type = {
+          'x': parsing_ops.FixedLenFeature([1], dtypes.float32, 33.0),
+          'y': parsing_ops.VarLenFeature(dtypes.string)
+      }
+      result = parsing_ops.parse_example(examples, feature_to_type)
+      parse_example_op = result['x'].op
+      config = extract_example_parser_configuration(parse_example_op, sess)
+      expected = self.getExpectedConfig(parse_example_op.type)
+      self.assertProtoEquals(expected, config)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/fast_module_type.cc

@@ -0,0 +1,415 @@
+/* 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 <Python.h>
+
+#include <array>
+#include <vector>
+
+// clang-format off
+// These headers must be at the top, before including Python.h header
+// Otherwise, we get C2039 on MSVC due to 'copysign'
+#include "absl/log/check.h"
+#include "pybind11/complex.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+// clang-format on
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/synchronization/mutex.h"
+
+namespace py = pybind11;
+
+#ifdef Py_GIL_DISABLED
+constexpr int PY_MODULE_TYPE_TP_BASIC_SIZE = 80;  // Under Python 3.13t
+#else                                             // Py_GIL_DISABLED
+constexpr int PY_MODULE_TYPE_TP_BASIC_SIZE = 56;
+#endif                                            // Py_GIL_DISABLED
+
+struct FastModuleObject {
+  // A dummy array that ensures enough size is reserved for FastModuleObject,
+  // because it's inherited from PyModuleObject.
+  const std::array<char, PY_MODULE_TYPE_TP_BASIC_SIZE> opaque_base_fields;
+  // A cache that helps reduce attribute lookup overhead.
+  absl::flat_hash_map<PyObject*, PyObject*> attr_map ABSL_GUARDED_BY(mutex);
+  // pointer to the external getattribute function
+  PyObject *cb_getattribute;
+  // pointer to the external getattr function
+  PyObject *cb_getattr;
+  // mutex to protect attr_map
+  absl::Mutex mutex;
+  // static PyTypeObject type;
+
+  FastModuleObject() = delete;
+  ~FastModuleObject() = delete;
+  static FastModuleObject *UncheckedCast(PyObject *obj);
+};
+
+static PyObject *FastModule_new(PyTypeObject *subtype, PyObject *args,
+                                PyObject *kwds) {
+  DCHECK_EQ(PY_MODULE_TYPE_TP_BASIC_SIZE, PyModule_Type.tp_basicsize);
+  PyObject *obj = PyModule_Type.tp_new(subtype, args, kwds);
+  FastModuleObject *self = reinterpret_cast<FastModuleObject *>(obj);
+  new (&(self->attr_map)) absl::flat_hash_map<PyObject *, PyObject *>();
+  new (&(self->mutex)) absl::Mutex();
+  self->cb_getattribute = nullptr;
+  self->cb_getattr = nullptr;
+  return obj;
+}
+
+static int FastModule_traverse(PyObject *self, visitproc visit, void *arg) {
+  if (int super_result = PyModule_Type.tp_traverse(self, visit, arg) != 0) {
+    return super_result;
+  }
+  FastModuleObject* fast_module = FastModuleObject::UncheckedCast(self);
+  {
+    absl::MutexLock lock(&fast_module->mutex);
+    for (auto& it : fast_module->attr_map) {
+      Py_VISIT(it.first);
+      Py_VISIT(it.second);
+    }
+    // Visit the callbacks to ensure the GC tracks them and detects cycles.
+    Py_VISIT(fast_module->cb_getattribute);
+    Py_VISIT(fast_module->cb_getattr);
+  }
+  return 0;
+}
+
+// Parses the input as a callable and checks the result.
+static PyObject *ParseFunc(PyObject *args) {
+  PyObject *func;
+  if (!PyArg_ParseTuple(args, "O:set_callback", &func)) return nullptr;
+  if (!PyCallable_Check(func)) {
+    PyErr_SetString(PyExc_TypeError, "input args must be callable");
+    return nullptr;
+  }
+  Py_INCREF(func);  // Add a reference to new callback
+  return func;
+}
+
+// Sets the pointer 'cb_getattribute' in the FastModuleObject object
+// corresponding to 'self'.
+static PyObject *SetGetattributeCallback(PyObject *self, PyObject *args) {
+  PyObject *func = ParseFunc(args);
+  if (func == nullptr) return nullptr;
+  FastModuleObject* fast_module = FastModuleObject::UncheckedCast(self);
+  PyObject* old_func = nullptr;
+  {
+    absl::MutexLock lock(&fast_module->mutex);
+    old_func = fast_module->cb_getattribute;
+    fast_module->cb_getattribute = func;
+  }
+  Py_XDECREF(old_func);
+  Py_RETURN_NONE;
+}
+
+// Sets the pointer 'cb_getattr' in the FastModuleObject object
+// corresponding to 'self'.
+static PyObject *SetGetattrCallback(PyObject *self, PyObject *args) {
+  PyObject *func = ParseFunc(args);
+  if (func == nullptr) return nullptr;
+  FastModuleObject* fast_module = FastModuleObject::UncheckedCast(self);
+  PyObject* old_func = nullptr;
+  {
+    absl::MutexLock lock(&fast_module->mutex);
+    old_func = fast_module->cb_getattr;
+    fast_module->cb_getattr = func;
+  }
+  Py_XDECREF(old_func);
+  Py_RETURN_NONE;
+}
+
+// Inserts or updates a key-value pair in the cache 'attr_map'
+// of the FastModuleObject object corresponding to 'self'.
+static PyObject *FastDictInsert(FastModuleObject *self, PyObject *args) {
+  PyObject *name, *value;
+  if (!PyArg_ParseTuple(args, "OO", &name, &value)) {
+    PyErr_SetString(PyExc_TypeError, "_fastdict_insert: incorrect inputs");
+    return nullptr;
+  }
+  Py_INCREF(value);
+  PyObject* to_decref = nullptr;
+  {
+    absl::MutexLock lock(&self->mutex);
+    auto& attr_map = self->attr_map;
+    if (auto [it, inserted] = attr_map.emplace(name, value); inserted) {
+      Py_INCREF(name);
+    } else {
+      to_decref = it->second;
+      it->second = value;
+    }
+  }
+  Py_XDECREF(to_decref);
+
+  // Properly handle returning Py_None
+  Py_RETURN_NONE;
+}
+
+// Gets a value from a key in the cache 'attr_map'
+// of the FastModuleObject object corresponding to 'self'.
+static PyObject *FastDictGet(FastModuleObject *self, PyObject *args) {
+  PyObject *name;
+  if (!PyArg_ParseTuple(args, "O", &name)) {
+    PyErr_SetString(PyExc_TypeError, "_fastdict_get: incorrect inputs");
+    return nullptr;
+  }
+  {
+    absl::MutexLock lock(&self->mutex);
+    auto& attr_map = self->attr_map;
+    if (auto it = attr_map.find(name); it != attr_map.end()) {
+      Py_INCREF(it->second);
+      return it->second;
+    }
+  }
+  // Copied from CPython's moduleobject.c
+  PyErr_Format(PyExc_KeyError, "module has no attribute '%U'", name);
+  return nullptr;
+}
+
+// Gets a value from a key in the cache 'attr_map'
+// of the FastModuleObject object corresponding to 'self'.
+static PyObject *FastDictPop(FastModuleObject *self, PyObject *args) {
+  PyObject *name;
+  if (!PyArg_ParseTuple(args, "O", &name)) {
+    PyErr_SetString(PyExc_TypeError, "_fastdict_pop: incorrect inputs");
+    return nullptr;
+  }
+  PyObject* key_to_decref = nullptr;
+  PyObject* value = nullptr;
+  {
+    absl::MutexLock lock(&self->mutex);
+    auto& attr_map = self->attr_map;
+    if (auto it = attr_map.find(name); it != attr_map.end()) {
+      key_to_decref = it->first;
+      value = it->second;
+      attr_map.erase(it);
+    }
+  }
+  if (value) {
+    Py_DECREF(key_to_decref);
+    return value;
+  }
+  // Copied from CPython's moduleobject.c
+  PyErr_Format(PyExc_KeyError, "module has no attribute '%U'", name);
+  return nullptr;
+}
+
+// Returns true if a key exists in the cache 'attr_map'
+// of the FastModuleObject object corresponding to 'self',
+// otherwise returns false.
+static PyObject *FastDictContains(FastModuleObject *self, PyObject *args) {
+  PyObject *name;
+  if (!PyArg_ParseTuple(args, "O", &name)) {
+    PyErr_SetString(PyExc_TypeError, "_fastdict_key_in: incorrect inputs");
+    return nullptr;
+  }
+  bool result;
+  {
+    absl::MutexLock lock(&self->mutex);
+    const auto& attr_map = self->attr_map;
+    result = attr_map.contains(name);
+  }
+  if (result) {
+    // Properly handle returning Py_True
+    Py_RETURN_TRUE;
+  }
+  // Properly handle returning Py_False
+  Py_RETURN_FALSE;
+}
+
+// Calls a function 'func' with inputs 'self' and 'args'.
+static PyObject *CallFunc(FastModuleObject *self, PyObject *args,
+                          PyObject *func) {
+  if (func == nullptr) {
+    PyErr_SetString(PyExc_NameError,
+                    "Attempting to call a callback that was not defined");
+    return nullptr;
+  }
+  PyObject *name;
+  if (!PyArg_ParseTuple(args, "O", &name)) {
+    PyErr_SetString(PyExc_TypeError, "CallFunc: incorrect inputs");
+    return nullptr;
+  }
+  PyObject *arglist = Py_BuildValue("(OO)", self, name);
+  auto result = PyObject_CallObject(func, arglist);
+  Py_DECREF(arglist);
+  return result;
+}
+
+static PyMethodDef FastModule_methods[] = {
+    {"_fastdict_insert", reinterpret_cast<PyCFunction>(FastDictInsert),
+     METH_VARARGS, "Registers a method to the fast lookup table."},
+    {"_fastdict_get", reinterpret_cast<PyCFunction>(FastDictGet), METH_VARARGS,
+     "Gets a method from the fast lookup table."},
+    {"_fastdict_pop", reinterpret_cast<PyCFunction>(FastDictPop), METH_VARARGS,
+     "Removes a method in the fast lookup table."},
+    {"_fastdict_key_in", reinterpret_cast<PyCFunction>(FastDictContains),
+     METH_VARARGS, "Checks if a method exists in the fast lookup table."},
+    {"set_getattribute_callback", SetGetattributeCallback, METH_VARARGS,
+     "Defines the callback function to replace __getattribute__"},
+    {"set_getattr_callback", SetGetattrCallback, METH_VARARGS,
+     "Defines the callback function to replace __getattr__"},
+    {nullptr, nullptr, 0, nullptr},
+};
+
+// Attempts to get the attribute based on 'name' as the key in cache 'attr_map'
+// of the FastModuleObject object corresponding to 'module'.
+// If the lookup fails in the cache, either uses
+// a user-defined callback 'cb_getattribute'
+// or the default 'tp_getattro' function to look for the attribute.
+static PyObject *FastTpGetattro(PyObject *module, PyObject *name) {
+  FastModuleObject *fast_module = FastModuleObject::UncheckedCast(module);
+  {
+    absl::MutexLock lock(&fast_module->mutex);
+    auto& attr_map = fast_module->attr_map;
+    // If the attribute lookup is successful in the cache, directly return it.
+    if (auto it = attr_map.find(name); it != attr_map.end()) {
+      PyObject* value = it->second;
+      Py_INCREF(value);
+      return value;
+    }
+  }
+  PyObject *arglist = Py_BuildValue("(O)", name);
+  PyObject *result;
+  PyObject* cb_getattribute = nullptr;
+  {
+    absl::MutexLock lock(&fast_module->mutex);
+    cb_getattribute = fast_module->cb_getattribute;
+    Py_XINCREF(cb_getattribute);
+  }
+  // Prefer the customized callback function over the default function.
+  if (cb_getattribute != nullptr) {
+    result = CallFunc(fast_module, arglist, cb_getattribute);
+    Py_DECREF(cb_getattribute);
+  } else {
+    result = PyModule_Type.tp_getattro(module, name);
+  }
+  // Return result if it's found
+  if (result != nullptr) {
+    Py_DECREF(arglist);
+    return result;
+  }
+  // If the default lookup fails and an AttributeError is raised,
+  // clear the error status before using the __getattr__ callback function.
+  auto is_error = PyErr_Occurred();
+  PyObject* cb_getattr = nullptr;
+  if (is_error && PyErr_ExceptionMatches(PyExc_AttributeError)) {
+    absl::MutexLock lock(&fast_module->mutex);
+    cb_getattr = fast_module->cb_getattr;
+    Py_XINCREF(cb_getattr);
+  }
+  if (cb_getattr != nullptr) {
+    PyErr_Clear();
+    result = CallFunc(fast_module, arglist, cb_getattr);
+    Py_DECREF(cb_getattr);
+  }
+  // If all options were used up
+  Py_DECREF(arglist);
+  return result;
+}
+
+// Customized destructor for FastModuleType.tp_dealloc
+// In addition to default behavior it also clears up the contents in attr_map.
+static void FastModuleObjectDealloc(PyObject *module) {
+  FastModuleObject *fast_module = FastModuleObject::UncheckedCast(module);
+  std::vector<std::pair<PyObject*, PyObject*>> items;
+  {
+    absl::MutexLock lock(&fast_module->mutex);
+    items.reserve(fast_module->attr_map.size());
+    for (auto [key, value] : fast_module->attr_map) {
+      items.push_back({key, value});
+    }
+    fast_module->attr_map.~flat_hash_map<PyObject*, PyObject*>();
+  }
+  for (auto [key, value] : items) {
+    Py_DECREF(key);
+    Py_DECREF(value);
+  }
+  fast_module->mutex.~Mutex();
+  Py_XDECREF(fast_module->cb_getattribute);
+  Py_XDECREF(fast_module->cb_getattr);
+  Py_TYPE(module)->tp_free(module);
+}
+
+static PyTypeObject FastModuleType = []() {
+  PyTypeObject obj = {PyVarObject_HEAD_INIT(&PyType_Type, 0)};
+  obj.tp_name = "fast_module_type.FastModuleType";
+  obj.tp_basicsize = sizeof(FastModuleObject);
+  obj.tp_itemsize = 0;
+  obj.tp_dealloc = FastModuleObjectDealloc;
+  obj.tp_getattro = FastTpGetattro;
+  obj.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC;
+  obj.tp_doc = "FastModuleType objects";
+  obj.tp_methods = FastModule_methods;
+  obj.tp_traverse = reinterpret_cast<traverseproc>(FastModule_traverse);
+  obj.tp_new = reinterpret_cast<newfunc>(FastModule_new);
+  return obj;
+}();
+
+// Returns true if the type of 'obj' or any of its parent class
+// is equal to 'target'. Otherwise returns false.
+bool IsAnyBaseSameType(const PyObject *obj, const PyTypeObject *target) {
+  auto *tp = Py_TYPE(obj);
+  while (true) {
+    if (tp == target) return true;
+    // If the default type is found, there is no need to search further
+    if (tp == &PyBaseObject_Type) break;
+    tp = tp->tp_base;
+  }
+  return false;
+}
+
+// Casts 'obj' to 'FastModuleObject *'.
+// Conducts a check only in non-optimized builds.
+FastModuleObject *FastModuleObject::UncheckedCast(PyObject *obj) {
+  DCHECK(IsAnyBaseSameType(obj, &FastModuleType));
+  return reinterpret_cast<FastModuleObject *>(obj);
+}
+
+PYBIND11_MODULE(fast_module_type, m, pybind11::mod_gil_not_used()) {
+  FastModuleType.tp_base = &PyModule_Type;
+  FastModuleType.tp_setattro = [](PyObject* module, PyObject* name,
+                                  PyObject* value) -> int {
+    auto* fast_module = FastModuleObject::UncheckedCast(module);
+    PyObject* to_decref = nullptr;
+    Py_INCREF(value);
+    {
+      absl::MutexLock lock(&fast_module->mutex);
+      auto& attr_map = fast_module->attr_map;
+      if (auto [it, inserted] = attr_map.emplace(name, value); inserted) {
+        Py_INCREF(name);
+      } else {
+        to_decref = it->second;
+        it->second = value;
+      }
+    }
+    Py_XDECREF(to_decref);
+    PyObject_GenericSetAttr(module, name, value);
+    return 0;
+  };
+
+  m.doc() = R"pbdoc(
+    fast_module_type
+    -----
+  )pbdoc";
+  // Use getter function to hold attributes rather than pybind11's m.attr due to
+  // b/145559202.
+  m.def(
+      "get_fast_module_type_class",
+      []() {
+        return py::cast<py::object>(
+            reinterpret_cast<PyObject *>(&FastModuleType));
+      },
+      py::return_value_policy::reference);
+}

tensorflow/python/util/fast_module_type.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_fast_module_type_class() -> object: ...

tensorflow/python/util/fast_module_type_test.py

@@ -0,0 +1,107 @@
+# 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 tensorflow.python.util.fast_module_type."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import fast_module_type
+FastModuleType = fast_module_type.get_fast_module_type_class()
+
+
+class ChildFastModule(FastModuleType):
+
+  def _getattribute1(self, name):  # pylint: disable=unused-argument
+    return 2
+
+  def _getattribute2(self, name):  # pylint: disable=unused-argument
+    raise AttributeError("Pass to getattr")
+
+  def _getattr(self, name):  # pylint: disable=unused-argument
+    return 3
+
+
+class EarlyAttrAccessModule(FastModuleType):
+
+  def __init__(self, name):
+    self.some_attr = 1
+    super().__init__(name)  # pytype: disable=too-many-function-args
+
+
+class FastModuleTypeTest(test.TestCase):
+
+  def testAttributeAccessBeforeSuperInitDoesNotCrash(self):
+    # Tests that the attribute access before super().__init__() does not crash.
+    module = EarlyAttrAccessModule("early_attr")
+    self.assertEqual(1, module.some_attr)
+
+  def testMissingModuleNameCallDoesNotCrash(self):
+    with self.assertRaises(TypeError):
+      ChildFastModule()
+
+  def testBaseGetattribute(self):
+    # Tests that the default attribute lookup works.
+    module = ChildFastModule("test")
+    module.foo = 1
+    self.assertEqual(1, module.foo)
+
+  def testGetattributeCallback(self):
+    # Tests that functionality of __getattribute__ can be set as a callback.
+    module = ChildFastModule("test")
+    FastModuleType.set_getattribute_callback(module,
+                                             ChildFastModule._getattribute1)
+    self.assertEqual(2, module.foo)
+
+  def testInvalidGetattributeCallbackRaisesAndKeepsExistingCallback(self):
+    module = ChildFastModule("test")
+    FastModuleType.set_getattribute_callback(
+        module, ChildFastModule._getattribute1
+    )
+
+    with self.assertRaisesRegex(TypeError, "input args must be callable"):
+      FastModuleType.set_getattribute_callback(module, 1)
+
+    self.assertEqual(2, module.foo)
+
+  def testGetattrCallback(self):
+    # Tests that functionality of __getattr__ can be set as a callback.
+    module = ChildFastModule("test")
+    FastModuleType.set_getattribute_callback(module,
+                                             ChildFastModule._getattribute2)
+    FastModuleType.set_getattr_callback(module, ChildFastModule._getattr)
+    self.assertEqual(3, module.foo)
+
+  def testInvalidGetattrCallbackRaisesAndKeepsExistingCallback(self):
+    module = ChildFastModule("test")
+    FastModuleType.set_getattr_callback(module, ChildFastModule._getattr)
+
+    with self.assertRaisesRegex(TypeError, "input args must be callable"):
+      FastModuleType.set_getattr_callback(module, 1)
+
+    self.assertEqual(3, module.foo)
+
+  def testFastdictApis(self):
+    module = ChildFastModule("test")
+    # At first "bar" does not exist in the module's attributes
+    self.assertFalse(module._fastdict_key_in("bar"))
+    with self.assertRaisesRegex(KeyError, "module has no attribute 'bar'"):
+      module._fastdict_get("bar")
+
+    module._fastdict_insert("bar", 1)
+    # After _fastdict_insert() the attribute is added.
+    self.assertTrue(module._fastdict_key_in("bar"))
+    self.assertEqual(1, module.bar)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/function_parameter_canonicalizer.cc

@@ -0,0 +1,176 @@
+/* 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/python/util/function_parameter_canonicalizer.h"
+
+#include "absl/base/attributes.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/python/lib/core/py_util.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+
+namespace {
+inline const char* PyUnicodeAsUtf8Compat(PyObject* obj) {
+#if PY_MAJOR_VERSION < 3
+  return PyString_AS_STRING(obj);
+#else
+  return PyUnicode_AsUTF8(obj);
+#endif
+}
+
+inline PyObject* PyUnicodeInternFromStringCompat(const char* str) {
+#if PY_MAJOR_VERSION < 3
+  return PyString_InternFromString(str);
+#else
+  return PyUnicode_InternFromString(str);
+#endif
+}
+
+inline void PyUnicodeInternInPlaceCompat(PyObject** obj) {
+#if PY_MAJOR_VERSION < 3
+  PyString_InternInPlace(obj);
+#else
+  PyUnicode_InternInPlace(obj);
+#endif
+}
+
+}  // namespace
+
+namespace tensorflow {
+
+FunctionParameterCanonicalizer::FunctionParameterCanonicalizer(
+    absl::Span<const char*> arg_names, absl::Span<PyObject*> defaults)
+    : positional_args_size_(arg_names.size() - defaults.size()) {
+  DCheckPyGilState();
+  DCHECK_GE(positional_args_size_, 0);
+
+  interned_arg_names_.reserve(arg_names.size());
+  for (const char* obj : arg_names)
+    interned_arg_names_.emplace_back(PyUnicodeInternFromStringCompat(obj));
+
+  DCHECK(AreInternedArgNamesUnique());
+
+  for (PyObject* obj : defaults) Py_INCREF(obj);
+  defaults_ = std::vector<Safe_PyObjectPtr>(defaults.begin(), defaults.end());
+}
+
+bool FunctionParameterCanonicalizer::Canonicalize(
+    PyObject* args, PyObject* kwargs, absl::Span<PyObject*> result) {
+  // TODO(kkb): Closely follow `Python/ceval.c`'s logic and error handling.
+
+  DCheckPyGilState();
+  DCHECK(PyTuple_CheckExact(args));
+  DCHECK(kwargs == nullptr || PyDict_CheckExact(kwargs));
+  DCHECK_EQ(result.size(), interned_arg_names_.size());
+
+  const int args_size = Py_SIZE(args);
+  int remaining_positional_args_count = positional_args_size_ - args_size;
+
+  // Check if the number of input arguments are too many.
+  if (TF_PREDICT_FALSE(args_size > interned_arg_names_.size())) {
+    PyErr_SetString(
+        PyExc_TypeError,
+        absl::StrCat("Too many arguments were given. Expected ",
+                     interned_arg_names_.size(), " but got ", args_size, ".")
+            .c_str());
+    return false;
+  }
+
+  // Fill positional arguments.
+  for (int i = 0; i < args_size; ++i) result[i] = PyTuple_GET_ITEM(args, i);
+
+  // Fill default arguments.
+  for (int i = std::max(positional_args_size_, args_size);
+       i < interned_arg_names_.size(); ++i)
+    result[i] = defaults_[i - positional_args_size_].get();
+
+  // Fill keyword arguments.
+  if (kwargs != nullptr) {
+    PyObject *key, *value;
+    Py_ssize_t pos = 0;
+    while (PyDict_Next(kwargs, &pos, &key, &value)) {
+      std::size_t index = InternedArgNameLinearSearch(key);
+
+      // Check if key object(argument name) was found in the pre-built intern
+      // string table.
+      if (TF_PREDICT_FALSE(index == interned_arg_names_.size())) {
+        // `key` might not be an interend string, so get the interned string
+        // and try again.  Note: we need to call INCREF before we use
+        // InternInPlace, to prevent the key in the dictionary from being
+        // prematurely deleted in the case where InternInPlace switches `key`
+        // to point at a new object.  We call DECREF(key) once we're done
+        // (which might decref the original key *or* the interned version).
+        Py_INCREF(key);
+        PyUnicodeInternInPlaceCompat(&key);
+        index = InternedArgNameLinearSearch(key);
+        Py_DECREF(key);
+
+        // Still not found, then return an error.
+        if (TF_PREDICT_FALSE(index == interned_arg_names_.size())) {
+          PyErr_Format(PyExc_TypeError,
+                       "Got an unexpected keyword argument '%s'",
+                       PyUnicodeAsUtf8Compat(key));
+          return false;
+        }
+      }
+
+      // Check if the keyword argument overlaps with positional arguments.
+      if (TF_PREDICT_FALSE(index < args_size)) {
+        PyErr_Format(PyExc_TypeError, "Got multiple values for argument '%s'",
+                     PyUnicodeAsUtf8Compat(key));
+        return false;
+      }
+
+      if (TF_PREDICT_FALSE(index < positional_args_size_))
+        --remaining_positional_args_count;
+
+      result[index] = value;
+    }
+  }
+
+  // Check if all the arguments are filled.
+  // Example failure, not enough number of arguments passed: `matmul(x)`
+  if (TF_PREDICT_FALSE(remaining_positional_args_count > 0)) {
+    // TODO(kkb): Report what arguments are missing.
+    PyErr_SetString(PyExc_TypeError, "Missing required positional argument");
+    return false;
+  }
+
+  return true;
+}
+
+ABSL_MUST_USE_RESULT
+ABSL_ATTRIBUTE_HOT
+inline std::size_t FunctionParameterCanonicalizer::InternedArgNameLinearSearch(
+    PyObject* name) {
+  std::size_t result = interned_arg_names_.size();
+
+  for (std::size_t i = 0; i < interned_arg_names_.size(); ++i)
+    if (TF_PREDICT_FALSE(name == interned_arg_names_[i].get())) return i;
+
+  return result;
+}
+
+bool FunctionParameterCanonicalizer::AreInternedArgNamesUnique() {
+  absl::flat_hash_set<PyObject*> interned_arg_names_set;
+  for (const Safe_PyObjectPtr& obj : interned_arg_names_)
+    interned_arg_names_set.emplace(obj.get());
+
+  return interned_arg_names_set.size() == interned_arg_names_.size();
+}
+}  // namespace tensorflow

tensorflow/python/util/function_parameter_canonicalizer.h

@@ -0,0 +1,74 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_PYTHON_UTIL_FUNCTION_PARAMETER_CANONICALIZER_H_
+#define TENSORFLOW_PYTHON_UTIL_FUNCTION_PARAMETER_CANONICALIZER_H_
+
+#include <Python.h>
+
+#include <cstddef>
+#include <vector>
+
+#include "absl/base/attributes.h"
+#include "absl/types/span.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+
+namespace tensorflow {
+
+// A class that Canonicalizes Python arg & kwargs parameters.
+class FunctionParameterCanonicalizer {
+ public:
+  // `arg_names` is a list of argument names, and `defaults` is default PyObject
+  // instances for arguments. `default` is aligned to the end.
+  FunctionParameterCanonicalizer(absl::Span<const char*> arg_names,
+                                 absl::Span<PyObject*> defaults);
+
+  // Returns the total number of arguments.
+  ABSL_MUST_USE_RESULT
+  int GetArgSize() const { return interned_arg_names_.size(); }
+
+  // Canonicalizes `args` and `kwargs` by the spec specified at construction.
+  // It's written to `result`. Returns `true` if Canonicalization was
+  // successful, and `false` otherwise. When it fails, it also sets CPython
+  // error status.
+  // This function does not update reference counter of any Python objects.
+  // `PyObject*`s in `result` are borrowed references from `args`, `kwargs`, and
+  // possibly `defaults_`, and will be only valid if `args` and `kwargs` are
+  // still alive.
+  ABSL_MUST_USE_RESULT
+  ABSL_ATTRIBUTE_HOT
+  bool Canonicalize(PyObject* args, PyObject* kwargs,
+                    absl::Span<PyObject*> result);
+
+ private:
+  // Simple linear search of `name` in `interned_arg_names`. If found, returns
+  // the index. If not found, returns `interned_arg_names.size()`.
+  ABSL_MUST_USE_RESULT
+  ABSL_ATTRIBUTE_HOT
+  std::size_t InternedArgNameLinearSearch(PyObject* name);
+
+  // Check if `interned_arg_names_` is unique.
+  bool AreInternedArgNamesUnique();
+
+  // TODO(kkb): Use one `std::vector` and two `absl:Span`s instead to improve
+  // cache locality.
+  std::vector<Safe_PyObjectPtr> interned_arg_names_;
+  std::vector<Safe_PyObjectPtr> defaults_;
+  const int positional_args_size_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_UTIL_FUNCTION_PARAMETER_CANONICALIZER_H_

tensorflow/python/util/function_parameter_canonicalizer_binding_for_test.cc

@@ -0,0 +1,75 @@
+/* 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 <Python.h>
+
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/types/span.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+#include "tensorflow/python/util/function_parameter_canonicalizer.h"
+
+namespace py = pybind11;
+
+class FunctionParameterCanonicalizerWrapper {
+ public:
+  FunctionParameterCanonicalizerWrapper(absl::Span<const char*> arg_names,
+                                        absl::Span<PyObject*> defaults)
+      : function_parameter_canonicalizer_(arg_names, defaults) {}
+
+  tensorflow::FunctionParameterCanonicalizer function_parameter_canonicalizer_;
+};
+
+PYBIND11_MODULE(_function_parameter_canonicalizer_binding_for_test, m) {
+  py::class_<FunctionParameterCanonicalizerWrapper>(
+      m, "FunctionParameterCanonicalizer")
+      .def(py::init([](std::vector<std::string> arg_names, py::tuple defaults) {
+        std::vector<const char*> arg_names_c_str;
+        for (const std::string& name : arg_names)
+          arg_names_c_str.emplace_back(name.c_str());
+
+        tensorflow::Safe_PyObjectPtr defaults_fast(
+            PySequence_Fast(defaults.ptr(), "Expected tuple"));
+        if (!defaults) throw py::error_already_set();
+        PyObject** default_items = PySequence_Fast_ITEMS(defaults_fast.get());
+        return new FunctionParameterCanonicalizerWrapper(
+            absl::MakeSpan(arg_names_c_str),
+            absl::MakeSpan(default_items,
+                           PySequence_Fast_GET_SIZE(defaults_fast.get())));
+      }))
+      .def("canonicalize", [](FunctionParameterCanonicalizerWrapper& self,
+                              py::args args, py::kwargs kwargs) {
+        std::vector<PyObject*> result_raw(
+            self.function_parameter_canonicalizer_.GetArgSize());
+
+        bool is_suceeded = self.function_parameter_canonicalizer_.Canonicalize(
+            args.ptr(), kwargs.ptr(), absl::MakeSpan(result_raw));
+
+        if (!is_suceeded) {
+          CHECK(PyErr_Occurred());
+          throw py::error_already_set();
+        }
+
+        py::list result;
+        for (PyObject* obj : result_raw) result.append(obj);
+        return result;
+      });
+}

tensorflow/python/util/function_parameter_canonicalizer_test.py

@@ -0,0 +1,93 @@
+# 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::FunctionParameterCanonicalizer`."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import _function_parameter_canonicalizer_binding_for_test
+
+
+class FunctionParameterCanonicalizerTest(test.TestCase):
+
+  def setUp(self):
+    super(FunctionParameterCanonicalizerTest, self).setUp()
+    self._matmul_func = (
+        _function_parameter_canonicalizer_binding_for_test
+        .FunctionParameterCanonicalizer([
+            'a', 'b', 'transpose_a', 'transpose_b', 'adjoint_a', 'adjoint_b',
+            'a_is_sparse', 'b_is_sparse', 'name'
+        ], (False, False, False, False, False, False, None)))
+
+  def testPosOnly(self):
+    self.assertEqual(
+        self._matmul_func.canonicalize(2, 3),
+        [2, 3, False, False, False, False, False, False, None])
+
+  def testPosOnly2(self):
+    self.assertEqual(
+        self._matmul_func.canonicalize(2, 3, True, False, True),
+        [2, 3, True, False, True, False, False, False, None])
+
+  def testPosAndKwd(self):
+    self.assertEqual(
+        self._matmul_func.canonicalize(
+            2, 3, transpose_a=True, name='my_matmul'),
+        [2, 3, True, False, False, False, False, False, 'my_matmul'])
+
+  def testPosAndKwd2(self):
+    self.assertEqual(
+        self._matmul_func.canonicalize(2, b=3),
+        [2, 3, False, False, False, False, False, False, None])
+
+  def testMissingPos(self):
+    with self.assertRaisesRegex(TypeError,
+                                'Missing required positional argument'):
+      self._matmul_func.canonicalize(2)
+
+  def testMissingPos2(self):
+    with self.assertRaisesRegex(TypeError,
+                                'Missing required positional argument'):
+      self._matmul_func.canonicalize(
+          transpose_a=True, transpose_b=True, adjoint_a=True)
+
+  def testTooManyArgs(self):
+    with self.assertRaisesRegex(TypeError, 'Too many arguments were given.'
+                                ' Expected 9 but got 10.'):
+      self._matmul_func.canonicalize(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
+
+  def testInvalidKwd(self):
+    with self.assertRaisesRegex(TypeError,
+                                'Got an unexpected keyword argument'):
+      self._matmul_func.canonicalize(2, 3, hohoho=True)
+
+  def testDuplicatedArg(self):
+    with self.assertRaisesRegex(TypeError,
+                                "Got multiple values for argument 'b'"):
+      self._matmul_func.canonicalize(2, 3, False, b=4)
+
+  def testDuplicatedArg2(self):
+    with self.assertRaisesRegex(
+        TypeError, "Got multiple values for argument 'transpose_a'"):
+      self._matmul_func.canonicalize(2, 3, False, transpose_a=True)
+
+  def testKwargNotInterned(self):
+    func = (
+        _function_parameter_canonicalizer_binding_for_test
+        .FunctionParameterCanonicalizer(['long_parameter_name'], ()))
+    kwargs = dict([('_'.join(['long', 'parameter', 'name']), 5)])
+    func.canonicalize(**kwargs)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/function_utils.py

@@ -0,0 +1,132 @@
+# 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 retrieve function args."""
+
+import functools
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+
+
+def _is_bound_method(fn):
+  _, fn = tf_decorator.unwrap(fn)
+  return tf_inspect.ismethod(fn) and (fn.__self__ is not None)
+
+
+def _is_callable_object(obj):
+  return hasattr(obj, '__call__') and tf_inspect.ismethod(obj.__call__)
+
+
+def fn_args(fn):
+  """Get argument names for function-like object.
+
+  Args:
+    fn: Function, or function-like object (e.g., result of `functools.partial`).
+
+  Returns:
+    `tuple` of string argument names.
+
+  Raises:
+    ValueError: if partial function has positionally bound arguments
+  """
+  if isinstance(fn, functools.partial):
+    args = fn_args(fn.func)
+    args = [a for a in args[len(fn.args):] if a not in (fn.keywords or [])]
+  else:
+    if _is_callable_object(fn):
+      fn = fn.__call__
+    args = tf_inspect.getfullargspec(fn).args
+    if _is_bound_method(fn) and args:
+      # If it's a bound method, it may or may not have a self/cls first
+      # argument; for example, self could be captured in *args.
+      # If it does have a positional argument, it is self/cls.
+      args.pop(0)
+  return tuple(args)
+
+
+def has_kwargs(fn):
+  """Returns whether the passed callable has **kwargs in its signature.
+
+  Args:
+    fn: Function, or function-like object (e.g., result of `functools.partial`).
+
+  Returns:
+    `bool`: if `fn` has **kwargs in its signature.
+
+  Raises:
+     `TypeError`: If fn is not a Function, or function-like object.
+  """
+  if isinstance(fn, functools.partial):
+    fn = fn.func
+  elif _is_callable_object(fn):
+    fn = fn.__call__
+  elif not callable(fn):
+    raise TypeError(
+        'Argument `fn` should be a callable. '
+        f'Received: fn={fn} (of type {type(fn)})')
+  return tf_inspect.getfullargspec(fn).varkw is not None
+
+
+def get_func_name(func):
+  """Returns name of passed callable."""
+  _, func = tf_decorator.unwrap(func)
+  if callable(func):
+    if tf_inspect.isfunction(func):
+      return func.__name__
+    elif tf_inspect.ismethod(func):
+      return '%s.%s' % (
+          func.__self__.__class__.__name__,
+          func.__func__.__name__,
+      )
+    else:  # Probably a class instance with __call__
+      return str(type(func))
+  else:
+    raise ValueError(
+        'Argument `func` must be a callable. '
+        f'Received func={func} (of type {type(func)})')
+
+
+def get_func_code(func):
+  """Returns func_code of passed callable, or None if not available."""
+  _, func = tf_decorator.unwrap(func)
+  if callable(func):
+    if tf_inspect.isfunction(func) or tf_inspect.ismethod(func):
+      return func.__code__
+    # Since the object is not a function or method, but is a callable, we will
+    # try to access the __call__method as a function.  This works with callable
+    # classes but fails with functool.partial objects despite their __call__
+    # attribute.
+    try:
+      return func.__call__.__code__
+    except AttributeError:
+      return None
+  else:
+    raise ValueError(
+        'Argument `func` must be a callable. '
+        f'Received func={func} (of type {type(func)})')
+
+
+_rewriter_config_optimizer_disabled = None
+
+
+def get_disabled_rewriter_config():
+  global _rewriter_config_optimizer_disabled
+  if _rewriter_config_optimizer_disabled is None:
+    config = config_pb2.ConfigProto()
+    rewriter_config = config.graph_options.rewrite_options
+    rewriter_config.disable_meta_optimizer = True
+    _rewriter_config_optimizer_disabled = config.SerializeToString()
+  return _rewriter_config_optimizer_disabled

tensorflow/python/util/function_utils_test.py

@@ -0,0 +1,305 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for function args retrieval utils."""
+
+import functools
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import function_utils
+
+
+def silly_example_function():
+  pass
+
+
+class SillyCallableClass(object):
+
+  def __call__(self):
+    pass
+
+
+class FnArgsTest(test.TestCase):
+
+  def test_simple_function(self):
+    def fn(a, b):
+      return a + b
+    self.assertEqual(('a', 'b'), function_utils.fn_args(fn))
+
+  def test_callable(self):
+
+    class Foo(object):
+
+      def __call__(self, a, b):
+        return a + b
+
+    self.assertEqual(('a', 'b'), function_utils.fn_args(Foo()))
+
+  def test_bound_method(self):
+
+    class Foo(object):
+
+      def bar(self, a, b):
+        return a + b
+
+    self.assertEqual(('a', 'b'), function_utils.fn_args(Foo().bar))
+
+  def test_bound_method_no_self(self):
+
+    class Foo(object):
+
+      def bar(*args):  # pylint:disable=no-method-argument
+        return args[1] + args[2]
+
+    self.assertEqual((), function_utils.fn_args(Foo().bar))
+
+  def test_partial_function(self):
+    expected_test_arg = 123
+
+    def fn(a, test_arg):
+      if test_arg != expected_test_arg:
+        return ValueError('partial fn does not work correctly')
+      return a
+
+    wrapped_fn = functools.partial(fn, test_arg=123)
+
+    self.assertEqual(('a',), function_utils.fn_args(wrapped_fn))
+
+  def test_partial_function_with_positional_args(self):
+    expected_test_arg = 123
+
+    def fn(test_arg, a):
+      if test_arg != expected_test_arg:
+        return ValueError('partial fn does not work correctly')
+      return a
+
+    wrapped_fn = functools.partial(fn, 123)
+
+    self.assertEqual(('a',), function_utils.fn_args(wrapped_fn))
+
+    self.assertEqual(3, wrapped_fn(3))
+    self.assertEqual(3, wrapped_fn(a=3))
+
+  def test_double_partial(self):
+    expected_test_arg1 = 123
+    expected_test_arg2 = 456
+
+    def fn(a, test_arg1, test_arg2):
+      if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+        return ValueError('partial does not work correctly')
+      return a
+
+    wrapped_fn = functools.partial(fn, test_arg2=456)
+    double_wrapped_fn = functools.partial(wrapped_fn, test_arg1=123)
+
+    self.assertEqual(('a',), function_utils.fn_args(double_wrapped_fn))
+
+  def test_double_partial_with_positional_args_in_outer_layer(self):
+    expected_test_arg1 = 123
+    expected_test_arg2 = 456
+
+    def fn(test_arg1, a, test_arg2):
+      if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+        return ValueError('partial fn does not work correctly')
+      return a
+
+    wrapped_fn = functools.partial(fn, test_arg2=456)
+    double_wrapped_fn = functools.partial(wrapped_fn, 123)
+
+    self.assertEqual(('a',), function_utils.fn_args(double_wrapped_fn))
+
+    self.assertEqual(3, double_wrapped_fn(3))  # pylint: disable=no-value-for-parameter
+    self.assertEqual(3, double_wrapped_fn(a=3))  # pylint: disable=no-value-for-parameter
+
+  def test_double_partial_with_positional_args_in_both_layers(self):
+    expected_test_arg1 = 123
+    expected_test_arg2 = 456
+
+    def fn(test_arg1, test_arg2, a):
+      if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+        return ValueError('partial fn does not work correctly')
+      return a
+
+    wrapped_fn = functools.partial(fn, 123)  # binds to test_arg1
+    double_wrapped_fn = functools.partial(wrapped_fn, 456)  # binds to test_arg2
+
+    self.assertEqual(('a',), function_utils.fn_args(double_wrapped_fn))
+
+    self.assertEqual(3, double_wrapped_fn(3))  # pylint: disable=no-value-for-parameter
+    self.assertEqual(3, double_wrapped_fn(a=3))  # pylint: disable=no-value-for-parameter
+
+
+class HasKwargsTest(test.TestCase):
+
+  def test_simple_function(self):
+
+    fn_has_kwargs = lambda **x: x
+    self.assertTrue(function_utils.has_kwargs(fn_has_kwargs))
+
+    fn_has_no_kwargs = lambda x: x
+    self.assertFalse(function_utils.has_kwargs(fn_has_no_kwargs))
+
+  def test_callable(self):
+
+    class FooHasKwargs(object):
+
+      def __call__(self, **x):
+        del x
+    self.assertTrue(function_utils.has_kwargs(FooHasKwargs()))
+
+    class FooHasNoKwargs(object):
+
+      def __call__(self, x):
+        del x
+    self.assertFalse(function_utils.has_kwargs(FooHasNoKwargs()))
+
+  def test_bound_method(self):
+
+    class FooHasKwargs(object):
+
+      def fn(self, **x):
+        del x
+    self.assertTrue(function_utils.has_kwargs(FooHasKwargs().fn))
+
+    class FooHasNoKwargs(object):
+
+      def fn(self, x):
+        del x
+    self.assertFalse(function_utils.has_kwargs(FooHasNoKwargs().fn))
+
+  def test_partial_function(self):
+    expected_test_arg = 123
+
+    def fn_has_kwargs(test_arg, **x):
+      if test_arg != expected_test_arg:
+        return ValueError('partial fn does not work correctly')
+      return x
+
+    wrapped_fn = functools.partial(fn_has_kwargs, test_arg=123)
+    self.assertTrue(function_utils.has_kwargs(wrapped_fn))
+    some_kwargs = dict(x=1, y=2, z=3)
+    self.assertEqual(wrapped_fn(**some_kwargs), some_kwargs)
+
+    def fn_has_no_kwargs(x, test_arg):
+      if test_arg != expected_test_arg:
+        return ValueError('partial fn does not work correctly')
+      return x
+
+    wrapped_fn = functools.partial(fn_has_no_kwargs, test_arg=123)
+    self.assertFalse(function_utils.has_kwargs(wrapped_fn))
+    some_arg = 1
+    self.assertEqual(wrapped_fn(some_arg), some_arg)
+
+  def test_double_partial(self):
+    expected_test_arg1 = 123
+    expected_test_arg2 = 456
+
+    def fn_has_kwargs(test_arg1, test_arg2, **x):
+      if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+        return ValueError('partial does not work correctly')
+      return x
+
+    wrapped_fn = functools.partial(fn_has_kwargs, test_arg2=456)
+    double_wrapped_fn = functools.partial(wrapped_fn, test_arg1=123)
+
+    self.assertTrue(function_utils.has_kwargs(double_wrapped_fn))
+    some_kwargs = dict(x=1, y=2, z=3)
+    self.assertEqual(double_wrapped_fn(**some_kwargs), some_kwargs)
+
+    def fn_has_no_kwargs(x, test_arg1, test_arg2):
+      if test_arg1 != expected_test_arg1 or test_arg2 != expected_test_arg2:
+        return ValueError('partial does not work correctly')
+      return x
+
+    wrapped_fn = functools.partial(fn_has_no_kwargs, test_arg2=456)
+    double_wrapped_fn = functools.partial(wrapped_fn, test_arg1=123)
+
+    self.assertFalse(function_utils.has_kwargs(double_wrapped_fn))
+    some_arg = 1
+    self.assertEqual(double_wrapped_fn(some_arg), some_arg)  # pylint: disable=no-value-for-parameter
+
+  def test_raises_type_error(self):
+    with self.assertRaisesRegex(TypeError,
+                                'should be a callable'):
+      function_utils.has_kwargs('not a function')
+
+
+class GetFuncNameTest(test.TestCase):
+
+  def testWithSimpleFunction(self):
+    self.assertEqual(
+        'silly_example_function',
+        function_utils.get_func_name(silly_example_function))
+
+  def testWithClassMethod(self):
+    self.assertEqual(
+        'GetFuncNameTest.testWithClassMethod',
+        function_utils.get_func_name(self.testWithClassMethod))
+
+  def testWithCallableClass(self):
+    callable_instance = SillyCallableClass()
+    self.assertRegex(
+        function_utils.get_func_name(callable_instance),
+        '<.*SillyCallableClass.*>')
+
+  def testWithFunctoolsPartial(self):
+    partial = functools.partial(silly_example_function)
+    self.assertRegex(
+        function_utils.get_func_name(partial), '<.*functools.partial.*>')
+
+  def testWithLambda(self):
+    anon_fn = lambda x: x
+    self.assertEqual('<lambda>', function_utils.get_func_name(anon_fn))
+
+  def testRaisesWithNonCallableObject(self):
+    with self.assertRaises(ValueError):
+      function_utils.get_func_name(None)
+
+
+class GetFuncCodeTest(test.TestCase):
+
+  def testWithSimpleFunction(self):
+    code = function_utils.get_func_code(silly_example_function)
+    self.assertIsNotNone(code)
+    self.assertRegex(code.co_filename, 'function_utils_test.py')
+
+  def testWithClassMethod(self):
+    code = function_utils.get_func_code(self.testWithClassMethod)
+    self.assertIsNotNone(code)
+    self.assertRegex(code.co_filename, 'function_utils_test.py')
+
+  def testWithCallableClass(self):
+    callable_instance = SillyCallableClass()
+    code = function_utils.get_func_code(callable_instance)
+    self.assertIsNotNone(code)
+    self.assertRegex(code.co_filename, 'function_utils_test.py')
+
+  def testWithLambda(self):
+    anon_fn = lambda x: x
+    code = function_utils.get_func_code(anon_fn)
+    self.assertIsNotNone(code)
+    self.assertRegex(code.co_filename, 'function_utils_test.py')
+
+  def testWithFunctoolsPartial(self):
+    partial = functools.partial(silly_example_function)
+    code = function_utils.get_func_code(partial)
+    self.assertIsNone(code)
+
+  def testRaisesWithNonCallableObject(self):
+    with self.assertRaises(ValueError):
+      function_utils.get_func_code(None)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/is_in_graph_mode.py

@@ -0,0 +1,18 @@
+# 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.
+# ==============================================================================
+"""A function that tells you if the program is running in graph mode."""
+# Call IS_IN_GRAPH_MODE() when you want to know whether the thread is in
+# graph mode.  By default, we always are.
+IS_IN_GRAPH_MODE = lambda: True

tensorflow/python/util/keras_deps.py

@@ -0,0 +1,80 @@
+# 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.
+# ==============================================================================
+
+"""Interface that provides access to Keras dependencies.
+
+This library is a common interface that contains Keras functions needed by
+TensorFlow and TensorFlow Lite and is required as per the dependency inversion
+principle (https://en.wikipedia.org/wiki/Dependency_inversion_principle). As per
+this principle, high-level modules (eg: TensorFlow and TensorFlow Lite) should
+not depend on low-level modules (eg: Keras) and instead both should depend on a
+common interface such as this file.
+"""
+
+
+from tensorflow.python.util.tf_export import tf_export
+
+_KERAS_CALL_CONTEXT_FUNCTION = None
+_KERAS_CLEAR_SESSION_FUNCTION = None
+_KERAS_GET_SESSION_FUNCTION = None
+_KERAS_LOAD_MODEL_FUNCTION = None
+
+# TODO(b/169898786): Use the Keras public API when TFLite moves out of TF
+
+
+# Register functions
+@tf_export('__internal__.register_call_context_function', v1=[])
+def register_call_context_function(func):
+  global _KERAS_CALL_CONTEXT_FUNCTION
+  _KERAS_CALL_CONTEXT_FUNCTION = func
+
+
+@tf_export('__internal__.register_clear_session_function', v1=[])
+def register_clear_session_function(func):
+  global _KERAS_CLEAR_SESSION_FUNCTION
+  _KERAS_CLEAR_SESSION_FUNCTION = func
+
+
+@tf_export('__internal__.register_get_session_function', v1=[])
+def register_get_session_function(func):
+  global _KERAS_GET_SESSION_FUNCTION
+  _KERAS_GET_SESSION_FUNCTION = func
+
+
+@tf_export('__internal__.register_load_model_function', v1=[])
+def register_load_model_function(func):
+  global _KERAS_LOAD_MODEL_FUNCTION
+  _KERAS_LOAD_MODEL_FUNCTION = func
+
+
+# Get functions
+def get_call_context_function():
+  global _KERAS_CALL_CONTEXT_FUNCTION
+  return _KERAS_CALL_CONTEXT_FUNCTION
+
+
+def get_clear_session_function():
+  global _KERAS_CLEAR_SESSION_FUNCTION
+  return _KERAS_CLEAR_SESSION_FUNCTION
+
+
+def get_get_session_function():
+  global _KERAS_GET_SESSION_FUNCTION
+  return _KERAS_GET_SESSION_FUNCTION
+
+
+def get_load_model_function():
+  global _KERAS_LOAD_MODEL_FUNCTION
+  return _KERAS_LOAD_MODEL_FUNCTION

tensorflow/python/util/kernel_registry.cc

@@ -0,0 +1,66 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/python/util/kernel_registry.h"
+
+#include <string>
+
+#include "absl/log/log.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+namespace tensorflow {
+namespace swig {
+
+std::string TryFindKernelClass(const std::string& serialized_node_def) {
+  tensorflow::NodeDef node_def;
+  if (!node_def.ParseFromString(serialized_node_def)) {
+    LOG(WARNING) << "Error parsing node_def";
+    return "";
+  }
+
+  const tensorflow::OpRegistrationData* op_reg_data;
+  auto status =
+      tensorflow::OpRegistry::Global()->LookUp(node_def.op(), &op_reg_data);
+  if (!status.ok()) {
+    LOG(WARNING) << "Op " << node_def.op() << " not found: " << status;
+    return "";
+  }
+  AddDefaultsToNodeDef(op_reg_data->op_def, &node_def);
+
+  tensorflow::DeviceNameUtils::ParsedName parsed_name;
+  if (!tensorflow::DeviceNameUtils::ParseFullName(node_def.device(),
+                                                  &parsed_name)) {
+    LOG(WARNING) << "Failed to parse device from node_def: "
+                 << node_def.ShortDebugString();
+    return "";
+  }
+  std::string class_name = "";
+  status = tensorflow::FindKernelDef(
+      tensorflow::DeviceType(parsed_name.type.c_str()), node_def,
+      nullptr /* kernel_def */, &class_name);
+  if (!status.ok()) {
+    LOG(WARNING) << "Op [" << node_def.op() << "]: " << status;
+  }
+  return class_name;
+}
+
+}  // namespace swig
+}  // namespace tensorflow

tensorflow/python/util/kernel_registry.h

@@ -0,0 +1,34 @@
+/* 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.
+==============================================================================*/
+
+// Functions for getting information about kernels registered in the binary.
+#ifndef TENSORFLOW_PYTHON_UTIL_KERNEL_REGISTRY_H_
+#define TENSORFLOW_PYTHON_UTIL_KERNEL_REGISTRY_H_
+
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace swig {
+
+// Returns the kernel class name required to execute <node_def> on the device
+// type of <node_def.device>, or an empty string if the kernel class is not
+// found or the device name is invalid.
+std::string TryFindKernelClass(const std::string& serialized_node_def);
+
+}  // namespace swig
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_UTIL_KERNEL_REGISTRY_H_

tensorflow/python/util/kernel_registry_wrapper.cc

@@ -0,0 +1,27 @@
+/* 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 <string>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/python/util/kernel_registry.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_kernel_registry, m, pybind11::mod_gil_not_used()) {
+  m.def("TryFindKernelClass", [](const std::string& serialized_node_def) {
+    return py::bytes(tensorflow::swig::TryFindKernelClass(serialized_node_def));
+  });
+}

tensorflow/python/util/keyword_args.py

@@ -0,0 +1,50 @@
+# 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.
+# ==============================================================================
+
+"""Keyword args functions."""
+
+import functools
+
+from tensorflow.python.util import decorator_utils
+
+
+def keyword_args_only(func):
+  """Decorator for marking specific function accepting keyword args only.
+
+  This decorator raises a `ValueError` if the input `func` is called with any
+  non-keyword args. This prevents the caller from providing the arguments in
+  wrong order.
+
+  Args:
+    func: The function or method needed to be decorated.
+
+  Returns:
+    Decorated function or method.
+
+  Raises:
+    ValueError: If `func` is not callable.
+  """
+
+  decorator_utils.validate_callable(func, "keyword_args_only")
+  @functools.wraps(func)
+  def new_func(*args, **kwargs):
+    """Keyword args only wrapper."""
+    if args:
+      raise ValueError(
+          f"The function {func.__name__} only accepts keyword arguments. "
+          "Do not pass positional arguments. Received the following positional "
+          f"arguments: {args}")
+    return func(**kwargs)
+  return new_func

tensorflow/python/util/keyword_args_test.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+"""Keyword args tests."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import keyword_args
+
+
+class KeywordArgsTest(test.TestCase):
+
+  def test_keyword_args_only(self):
+
+    def func_without_decorator(a, b):
+      return a + b
+
+    @keyword_args.keyword_args_only
+    def func_with_decorator(a, b):
+      return func_without_decorator(a, b)
+
+    self.assertEqual(3, func_without_decorator(1, 2))
+    self.assertEqual(3, func_without_decorator(a=1, b=2))
+    self.assertEqual(3, func_with_decorator(a=1, b=2))
+
+    # Providing non-keyword args should fail.
+    with self.assertRaisesRegex(
+        ValueError, "only accepts keyword arguments"):
+      self.assertEqual(3, func_with_decorator(1, 2))
+
+    # Partially providing keyword args should fail.
+    with self.assertRaisesRegex(
+        ValueError, "only accepts keyword arguments"):
+      self.assertEqual(3, func_with_decorator(1, b=2))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/lazy_loader.py

@@ -0,0 +1,224 @@
+# 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 LazyLoader class."""
+
+import importlib
+import os
+import types
+from tensorflow.python.platform import tf_logging as logging
+
+_TENSORFLOW_LAZY_LOADER_PREFIX = "_tfll"
+
+
+class LazyLoader(types.ModuleType):
+  """Lazily import a module, mainly to avoid pulling in large dependencies.
+
+  `contrib`, and `ffmpeg` are examples of modules that are large and not always
+  needed, and this allows them to only be loaded when they are used.
+  """
+
+  # The lint error here is incorrect.
+  def __init__(self, local_name, parent_module_globals, name, warning=None):
+    self._tfll_local_name = local_name
+    self._tfll_parent_module_globals = parent_module_globals
+    self._tfll_warning = warning
+
+    # These members allows doctest correctly process this module member without
+    # triggering self._load(). self._load() mutates parent_module_globals and
+    # triggers a dict mutated during iteration error from doctest.py.
+    # - for from_module()
+    super().__setattr__("__module__", name.rsplit(".", 1)[0])
+    # - for is_routine()
+    super().__setattr__("__wrapped__", None)
+
+    super().__init__(name)
+
+  def _load(self):
+    """Load the module and insert it into the parent's globals."""
+    # Import the target module and insert it into the parent's namespace
+    module = importlib.import_module(self.__name__)
+    self._tfll_parent_module_globals[self._tfll_local_name] = module
+
+    # Emit a warning if one was specified
+    if self._tfll_warning:
+      logging.warning(self._tfll_warning)
+      # Make sure to only warn once.
+      self._tfll_warning = None
+
+    # Update this object's dict so that if someone keeps a reference to the
+    #   LazyLoader, lookups are efficient (__getattr__ is only called on lookups
+    #   that fail).
+    self.__dict__.update(module.__dict__)
+
+    return module
+
+  def __getattr__(self, name):
+    module = self._load()
+    return getattr(module, name)
+
+  def __setattr__(self, name, value):
+    if name.startswith(_TENSORFLOW_LAZY_LOADER_PREFIX):
+      super().__setattr__(name, value)
+    else:
+      module = self._load()
+      setattr(module, name, value)
+      self.__dict__[name] = value
+      try:
+        # check if the module has __all__
+        if name not in self.__all__ and name != "__all__":
+          self.__all__.append(name)
+      except AttributeError:
+        pass
+
+  def __delattr__(self, name):
+    if name.startswith(_TENSORFLOW_LAZY_LOADER_PREFIX):
+      super().__delattr__(name)
+    else:
+      module = self._load()
+      delattr(module, name)
+      self.__dict__.pop(name)
+      try:
+        # check if the module has __all__
+        if name in self.__all__:
+          self.__all__.remove(name)
+      except AttributeError:
+        pass
+
+  def __repr__(self):
+    # Carefully to not trigger _load, since repr may be called in very
+    # sensitive places.
+    return f"<LazyLoader {self.__name__} as {self._tfll_local_name}>"
+
+  def __dir__(self):
+    module = self._load()
+    return dir(module)
+
+  def __reduce__(self):
+    return importlib.import_module, (self.__name__,)
+
+
+class KerasLazyLoader(LazyLoader):
+  """LazyLoader that handles routing to different Keras version."""
+
+  def __init__(  # pylint: disable=super-init-not-called
+      self, parent_module_globals, mode=None, submodule=None, name="keras"):
+    self._tfll_parent_module_globals = parent_module_globals
+    self._tfll_mode = mode
+    self._tfll_submodule = submodule
+    self._tfll_name = name
+    self._tfll_initialized = False
+
+  def _initialize(self):
+    """Resolve the Keras version to use and initialize the loader."""
+    self._tfll_initialized = True
+    package_name = None
+    keras_version = None
+    if os.environ.get("TF_USE_LEGACY_KERAS", None) in ("true", "True", "1"):
+      try:
+        import tf_keras  # pylint: disable=g-import-not-at-top,unused-import
+
+        keras_version = "tf_keras"
+        if self._tfll_mode == "v1":
+          package_name = "tf_keras.api._v1.keras"
+        else:
+          package_name = "tf_keras.api._v2.keras"
+      except ImportError:
+        logging.warning(
+            "Your environment has TF_USE_LEGACY_KERAS set to True, but you "
+            "do not have the tf_keras package installed. You must install it "
+            "in order to use the legacy tf.keras. Install it via: "
+            "`pip install tf_keras`"
+        )
+    else:
+      try:
+        import keras  # pylint: disable=g-import-not-at-top
+
+        if keras.__version__.startswith("3."):
+          # This is the Keras 3.x case.
+          keras_version = "keras_3"
+          package_name = "keras._tf_keras.keras"
+        else:
+          # This is the Keras 2.x case.
+          keras_version = "keras_2"
+          if self._tfll_mode == "v1":
+            package_name = "keras.api._v1.keras"
+          else:
+            package_name = "keras.api._v2.keras"
+      except ImportError:
+        raise ImportError(  # pylint: disable=raise-missing-from
+            "Keras cannot be imported. Check that it is installed."
+        )
+
+    self._tfll_keras_version = keras_version
+    if keras_version is not None:
+      if self._tfll_submodule is not None:
+        package_name += "." + self._tfll_submodule
+      super().__init__(
+          self._tfll_name, self._tfll_parent_module_globals, package_name
+      )
+    else:
+      raise ImportError(  # pylint: disable=raise-missing-from
+          "Keras cannot be imported. Check that it is installed."
+      )
+
+  def __getattr__(self, item):
+    if item in ("_tfll_mode", "_tfll_initialized", "_tfll_name"):
+      return super(types.ModuleType, self).__getattribute__(item)
+    if not self._tfll_initialized:
+      self._initialize()
+    if self._tfll_keras_version == "keras_3":
+      if (
+          self._tfll_mode == "v1"
+          and not self._tfll_submodule
+          and item.startswith("compat.v1.")
+      ):
+        raise AttributeError(
+            "`tf.compat.v1.keras` is not available with Keras 3. Keras 3 has "
+            "no support for TF 1 APIs. You can install the `tf_keras` package "
+            "as an alternative, and set the environment variable "
+            "`TF_USE_LEGACY_KERAS=True` to configure TensorFlow to route "
+            "`tf.compat.v1.keras` to `tf_keras`."
+        )
+      elif (
+          self._tfll_mode == "v2"
+          and not self._tfll_submodule
+          and item.startswith("compat.v2.")
+      ):
+        raise AttributeError(
+            "`tf.compat.v2.keras` is not available with Keras 3. Just use "
+            "`import keras` instead."
+        )
+      elif self._tfll_submodule and self._tfll_submodule.startswith(
+          "__internal__.legacy."
+      ):
+        raise AttributeError(
+            f"`{item}` is not available with Keras 3."
+        )
+    module = self._load()
+    return getattr(module, item)
+
+  def __repr__(self):
+    if self._tfll_initialized:
+      return (
+          f"<KerasLazyLoader ({self._tfll_keras_version}) "
+          f"{self.__name__} as {self._tfll_local_name} mode={self._tfll_mode}>"
+      )
+    return "<KerasLazyLoader>"
+
+  def __dir__(self):
+    if not self._tfll_initialized:
+      self._initialize()
+    return super().__dir__()

tensorflow/python/util/lazy_loader_test.py

@@ -0,0 +1,70 @@
+# 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.
+# ==============================================================================
+"""lazy loader tests."""
+
+# pylint: disable=unused-import
+import doctest
+import inspect
+import pickle
+import types
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import lazy_loader
+from tensorflow.python.util import tf_inspect
+
+
+class LazyLoaderTest(test.TestCase):
+
+  def testDocTestDoesNotLoad(self):
+    module = types.ModuleType("mytestmodule")
+    module.foo = lazy_loader.LazyLoader("foo", module.__dict__, "os.path")
+
+    self.assertIsInstance(module.foo, lazy_loader.LazyLoader)
+
+    finder = doctest.DocTestFinder()
+    finder.find(module)
+
+    self.assertIsInstance(module.foo, lazy_loader.LazyLoader)
+
+  @test.mock.patch.object(logging, "warning", autospec=True)
+  def testLazyLoaderMock(self, mock_warning):
+    name = LazyLoaderTest.__module__
+    lazy_loader_module = lazy_loader.LazyLoader(
+        "lazy_loader_module", globals(), name, warning="Test warning.")
+
+    self.assertEqual(0, mock_warning.call_count)
+    lazy_loader_module.foo = 0
+    self.assertEqual(1, mock_warning.call_count)
+    foo = lazy_loader_module.foo
+    self.assertEqual(1, mock_warning.call_count)
+
+    # Check that values stayed the same
+    self.assertEqual(lazy_loader_module.foo, foo)
+
+
+class PickleTest(test.TestCase):
+
+  def testPickleLazyLoader(self):
+    name = PickleTest.__module__  # Try to pickle current module.
+    lazy_loader_module = lazy_loader.LazyLoader(
+        "lazy_loader_module", globals(), name)
+    restored = pickle.loads(pickle.dumps(lazy_loader_module))
+    self.assertEqual(restored.__name__, name)
+    self.assertIsNotNone(restored.PickleTest)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/lock_util.py

@@ -0,0 +1,130 @@
+# 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.
+# ==============================================================================
+"""Locking related utils."""
+
+import threading
+
+
+class GroupLock(object):
+  """A lock to allow many members of a group to access a resource exclusively.
+
+  This lock provides a way to allow access to a resource by multiple threads
+  belonging to a logical group at the same time, while restricting access to
+  threads from all other groups. You can think of this as an extension of a
+  reader-writer lock, where you allow multiple writers at the same time. We
+  made it generic to support multiple groups instead of just two - readers and
+  writers.
+
+  Simple usage example with two groups accessing the same resource:
+
+  ```python
+  lock = GroupLock(num_groups=2)
+
+  # In a member of group 0:
+  with lock.group(0):
+    # do stuff, access the resource
+    # ...
+
+  # In a member of group 1:
+  with lock.group(1):
+    # do stuff, access the resource
+    # ...
+  ```
+
+  Using as a context manager with `.group(group_id)` is the easiest way. You
+  can also use the `acquire` and `release` method directly.
+  """
+
+  __slots__ = ["_ready", "_num_groups", "_group_member_counts"]
+
+  def __init__(self, num_groups=2):
+    """Initialize a group lock.
+
+    Args:
+      num_groups: The number of groups that will be accessing the resource under
+        consideration. Should be a positive number.
+
+    Returns:
+      A group lock that can then be used to synchronize code.
+
+    Raises:
+      ValueError: If num_groups is less than 1.
+    """
+    if num_groups < 1:
+      raise ValueError(
+          "Argument `num_groups` must be a positive integer. "
+          f"Received: num_groups={num_groups}")
+    self._ready = threading.Condition(threading.Lock())
+    self._num_groups = num_groups
+    self._group_member_counts = [0] * self._num_groups
+
+  def group(self, group_id):
+    """Enter a context where the lock is with group `group_id`.
+
+    Args:
+      group_id: The group for which to acquire and release the lock.
+
+    Returns:
+      A context manager which will acquire the lock for `group_id`.
+    """
+    self._validate_group_id(group_id)
+    return self._Context(self, group_id)
+
+  def acquire(self, group_id):
+    """Acquire the group lock for a specific group `group_id`."""
+    self._validate_group_id(group_id)
+
+    self._ready.acquire()
+    while self._another_group_active(group_id):
+      self._ready.wait()
+    self._group_member_counts[group_id] += 1
+    self._ready.release()
+
+  def release(self, group_id):
+    """Release the group lock for a specific group `group_id`."""
+    self._validate_group_id(group_id)
+
+    self._ready.acquire()
+    self._group_member_counts[group_id] -= 1
+    if self._group_member_counts[group_id] == 0:
+      self._ready.notify_all()
+    self._ready.release()
+
+  def _another_group_active(self, group_id):
+    return any(
+        c > 0 for g, c in enumerate(self._group_member_counts) if g != group_id)
+
+  def _validate_group_id(self, group_id):
+    if group_id < 0 or group_id >= self._num_groups:
+      raise ValueError(
+          "Argument `group_id` should verify `0 <= group_id < num_groups` "
+          f"(with `num_groups={self._num_groups}`). "
+          f"Received: group_id={group_id}")
+
+  class _Context(object):
+    """Context manager helper for `GroupLock`."""
+
+    __slots__ = ["_lock", "_group_id"]
+
+    def __init__(self, lock, group_id):
+      self._lock = lock
+      self._group_id = group_id
+
+    def __enter__(self):
+      self._lock.acquire(self._group_id)
+
+    def __exit__(self, type_arg, value_arg, traceback_arg):
+      del type_arg, value_arg, traceback_arg
+      self._lock.release(self._group_id)

tensorflow/python/util/lock_util_test.py

@@ -0,0 +1,59 @@
+# 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 lock_util."""
+
+import random
+import time
+
+from absl.testing import parameterized
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import lock_util
+
+
+class GroupLockTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(1, 2, 3, 5, 10)
+  def testGroups(self, num_groups):
+    lock = lock_util.GroupLock(num_groups)
+    num_threads = 10
+    finished = set()
+
+    def thread_fn(thread_id):
+      time.sleep(random.random() * 0.1)
+      group_id = thread_id % num_groups
+      with lock.group(group_id):
+        time.sleep(random.random() * 0.1)
+        self.assertGreater(lock._group_member_counts[group_id], 0)
+        for g, c in enumerate(lock._group_member_counts):
+          if g != group_id:
+            self.assertEqual(0, c)
+        finished.add(thread_id)
+
+    threads = [
+        self.checkedThread(target=thread_fn, args=(i,))
+        for i in range(num_threads)
+    ]
+
+    for i in range(num_threads):
+      threads[i].start()
+    for i in range(num_threads):
+      threads[i].join()
+
+    self.assertEqual(set(range(num_threads)), finished)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/module_wrapper.py

@@ -0,0 +1,283 @@
+# 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.
+# ==============================================================================
+"""Provides wrapper for TensorFlow modules."""
+
+import importlib
+
+from tensorflow.python.eager import monitoring
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import fast_module_type
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+from tensorflow.tools.compatibility import all_renames_v2
+
+FastModuleType = fast_module_type.get_fast_module_type_class()
+_PER_MODULE_WARNING_LIMIT = 1
+compat_v1_usage_gauge = monitoring.BoolGauge('/tensorflow/api/compat/v1',
+                                             'compat.v1 usage')
+
+
+def get_rename_v2(name):
+  if name not in all_renames_v2.symbol_renames:
+    return None
+  return all_renames_v2.symbol_renames[name]
+
+
+def _call_location():
+  """Extracts the caller filename and line number as a string.
+
+  Returns:
+    A string describing the caller source location.
+  """
+  frame = tf_inspect.currentframe()
+  assert frame.f_back.f_code.co_name == '_tfmw_add_deprecation_warning', (
+      'This function should be called directly from '
+      '_tfmw_add_deprecation_warning, as the caller is identified '
+      'heuristically by chopping off the top stack frames.')
+
+  # We want to get stack frame 3 frames up from current frame,
+  # i.e. above __getattr__, _tfmw_add_deprecation_warning,
+  # and _call_location calls.
+  for _ in range(3):
+    parent = frame.f_back
+    if parent is None:
+      break
+    frame = parent
+  return '{}:{}'.format(frame.f_code.co_filename, frame.f_lineno)
+
+
+def contains_deprecation_decorator(decorators):
+  return any(d.decorator_name == 'deprecated' for d in decorators)
+
+
+def has_deprecation_decorator(symbol):
+  """Checks if given object has a deprecation decorator.
+
+  We check if deprecation decorator is in decorators as well as
+  whether symbol is a class whose __init__ method has a deprecation
+  decorator.
+  Args:
+    symbol: Python object.
+
+  Returns:
+    True if symbol has deprecation decorator.
+  """
+  decorators, symbol = tf_decorator.unwrap(symbol)
+  if contains_deprecation_decorator(decorators):
+    return True
+  if tf_inspect.isfunction(symbol):
+    return False
+  if not tf_inspect.isclass(symbol):
+    return False
+  if not hasattr(symbol, '__init__'):
+    return False
+  init_decorators, _ = tf_decorator.unwrap(symbol.__init__)
+  return contains_deprecation_decorator(init_decorators)
+
+
+class TFModuleWrapper(FastModuleType):
+  """Wrapper for TF modules to support deprecation messages and lazyloading."""
+  # Ensures that compat.v1 API usage is recorded at most once
+  compat_v1_usage_recorded = False
+
+  def __init__(
+      self,
+      wrapped,
+      module_name,
+      public_apis=None,
+      deprecation=True,
+      has_lite=False):
+    super(TFModuleWrapper, self).__init__(wrapped.__name__)
+    FastModuleType.set_getattr_callback(self, TFModuleWrapper._getattr)
+    FastModuleType.set_getattribute_callback(self,
+                                             TFModuleWrapper._getattribute)
+    self.__dict__.update(wrapped.__dict__)
+    # Prefix all local attributes with _tfmw_ so that we can
+    # handle them differently in attribute access methods.
+    self._tfmw_wrapped_module = wrapped
+    self._tfmw_module_name = module_name
+    self._tfmw_public_apis = public_apis
+    self._tfmw_print_deprecation_warnings = deprecation
+    self._tfmw_has_lite = has_lite
+    self._tfmw_is_compat_v1 = (wrapped.__name__.endswith('.compat.v1'))
+    # Set __all__ so that import * work for lazy loaded modules
+    if self._tfmw_public_apis:
+      self._tfmw_wrapped_module.__all__ = list(self._tfmw_public_apis.keys())
+      self.__all__ = list(self._tfmw_public_apis.keys())
+    else:
+      if hasattr(self._tfmw_wrapped_module, '__all__'):
+        self.__all__ = self._tfmw_wrapped_module.__all__
+      else:
+        self._tfmw_wrapped_module.__all__ = [
+            attr for attr in dir(self._tfmw_wrapped_module)
+            if not attr.startswith('_')
+        ]
+        self.__all__ = self._tfmw_wrapped_module.__all__
+
+    # names we already checked for deprecation
+    self._tfmw_deprecated_checked = set()
+    self._tfmw_warning_count = 0
+
+  def _tfmw_add_deprecation_warning(self, name, attr):
+    """Print deprecation warning for attr with given name if necessary."""
+    if (self._tfmw_warning_count < _PER_MODULE_WARNING_LIMIT and
+        name not in self._tfmw_deprecated_checked):
+
+      self._tfmw_deprecated_checked.add(name)
+
+      if self._tfmw_module_name:
+        full_name = 'tf.%s.%s' % (self._tfmw_module_name, name)
+      else:
+        full_name = 'tf.%s' % name
+      rename = get_rename_v2(full_name)
+      if rename and not has_deprecation_decorator(attr):
+        call_location = _call_location()
+        # skip locations in Python source
+        if not call_location.startswith('<'):
+          logging.warning(
+              'From %s: The name %s is deprecated. Please use %s instead.\n',
+              _call_location(), full_name, rename)
+          self._tfmw_warning_count += 1
+          return True
+    return False
+
+  def _tfmw_import_module(self, name):
+    """Lazily loading the modules."""
+    # We ignore 'app' because it is accessed in __init__.py of tf.compat.v1.
+    # That way, if a user only imports tensorflow.compat.v1, it is not
+    # considered v1 API usage.
+    if (self._tfmw_is_compat_v1 and name != 'app' and
+        not TFModuleWrapper.compat_v1_usage_recorded):
+      TFModuleWrapper.compat_v1_usage_recorded = True
+      compat_v1_usage_gauge.get_cell().set(True)
+
+    symbol_loc_info = self._tfmw_public_apis[name]
+    if symbol_loc_info[0]:
+      module = importlib.import_module(symbol_loc_info[0])
+      attr = getattr(module, symbol_loc_info[1])
+    else:
+      attr = importlib.import_module(symbol_loc_info[1])
+    setattr(self._tfmw_wrapped_module, name, attr)
+    self.__dict__[name] = attr
+    # Cache the pair
+    self._fastdict_insert(name, attr)
+    return attr
+
+  def _getattribute(self, name):
+    # pylint: disable=g-doc-return-or-yield,g-doc-args
+    """Imports and caches pre-defined API.
+
+    Warns if necessary.
+
+    This method is a replacement for __getattribute__(). It will be added into
+    the extended python module as a callback to reduce API overhead.
+    """
+    # Avoid infinite recursions
+    func__fastdict_insert = object.__getattribute__(self, '_fastdict_insert')
+
+    # Make sure we do not import from tensorflow/lite/__init__.py
+    if name == 'lite':
+      if self._tfmw_has_lite:
+        attr = self._tfmw_import_module(name)
+        setattr(self._tfmw_wrapped_module, 'lite', attr)
+        func__fastdict_insert(name, attr)
+        return attr
+  # Placeholder for Google-internal contrib error
+
+    attr = object.__getattribute__(self, name)
+
+    # Return and cache dunders and our own members.
+    # This is necessary to guarantee successful construction.
+    # In addition, all the accessed attributes used during the construction must
+    # begin with "__" or "_tfmw" or "_fastdict_".
+    if name.startswith('__') or name.startswith('_tfmw_') or name.startswith(
+        '_fastdict_'):
+      func__fastdict_insert(name, attr)
+      return attr
+
+    # Print deprecations, only cache functions after deprecation warnings have
+    # stopped.
+    if not (self._tfmw_print_deprecation_warnings and
+            self._tfmw_add_deprecation_warning(name, attr)):
+      func__fastdict_insert(name, attr)
+
+    return attr
+
+  def _getattr(self, name):
+    # pylint: disable=g-doc-return-or-yield,g-doc-args
+    """Imports and caches pre-defined API.
+
+    Warns if necessary.
+
+    This method is a replacement for __getattr__(). It will be added into the
+    extended python module as a callback to reduce API overhead. Instead of
+    relying on implicit AttributeError handling, this added callback function
+    will
+    be called explicitly from the extended C API if the default attribute lookup
+    fails.
+    """
+    try:
+      attr = getattr(self._tfmw_wrapped_module, name)
+    except AttributeError:
+    # Placeholder for Google-internal contrib error
+
+      if not self._tfmw_public_apis:
+        raise
+      if name not in self._tfmw_public_apis:
+        raise
+      attr = self._tfmw_import_module(name)
+
+    if self._tfmw_print_deprecation_warnings:
+      self._tfmw_add_deprecation_warning(name, attr)
+    return attr
+
+  def __setattr__(self, arg, val):
+    if not arg.startswith('_tfmw_'):
+      setattr(self._tfmw_wrapped_module, arg, val)
+      self.__dict__[arg] = val
+      if arg not in self.__all__ and arg != '__all__':
+        self.__all__.append(arg)
+      # Update the cache
+      if self._fastdict_key_in(arg):
+        self._fastdict_insert(arg, val)
+    super(TFModuleWrapper, self).__setattr__(arg, val)
+
+  def __dir__(self):
+    if self._tfmw_public_apis:
+      return list(
+          set(self._tfmw_public_apis.keys()).union(
+              set([
+                  attr for attr in dir(self._tfmw_wrapped_module)
+                  if not attr.startswith('_')
+              ])))
+    else:
+      return dir(self._tfmw_wrapped_module)
+
+  def __delattr__(self, name):
+    if name.startswith('_tfmw_'):
+      super(TFModuleWrapper, self).__delattr__(name)
+    else:
+      delattr(self._tfmw_wrapped_module, name)
+      self.__dict__.pop(name)
+      if name in self.__all__:
+        self.__all__.remove(name)
+      self._fastdict_pop(name)
+      # delattr(self._tfmw_wrapped_module, name)
+
+  def __repr__(self):
+    return self._tfmw_wrapped_module.__repr__()
+
+  def __reduce__(self):
+    return importlib.import_module, (self.__name__,)

tensorflow/python/util/module_wrapper_test.py

@@ -0,0 +1,197 @@
+# 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 tensorflow.python.util.module_wrapper."""
+
+import pickle
+import types
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import module_wrapper
+from tensorflow.python.util import tf_inspect
+from tensorflow.tools.compatibility import all_renames_v2
+
+module_wrapper._PER_MODULE_WARNING_LIMIT = 5
+
+
+class MockModule(types.ModuleType):
+  pass
+
+
+class DeprecationWrapperTest(test.TestCase):
+
+  def testWrapperIsAModule(self):
+    module = MockModule('test')
+    wrapped_module = module_wrapper.TFModuleWrapper(module, 'test')
+    self.assertTrue(tf_inspect.ismodule(wrapped_module))
+
+  @test.mock.patch.object(logging, 'warning', autospec=True)
+  def testDeprecationWarnings(self, mock_warning):
+    module = MockModule('test')
+    module.foo = 1
+    module.bar = 2
+    module.baz = 3
+    all_renames_v2.symbol_renames['tf.test.bar'] = 'tf.bar2'
+    all_renames_v2.symbol_renames['tf.test.baz'] = 'tf.compat.v1.baz'
+
+    wrapped_module = module_wrapper.TFModuleWrapper(module, 'test')
+    self.assertTrue(tf_inspect.ismodule(wrapped_module))
+
+    self.assertEqual(0, mock_warning.call_count)
+    bar = wrapped_module.bar
+    self.assertEqual(1, mock_warning.call_count)
+    foo = wrapped_module.foo
+    self.assertEqual(1, mock_warning.call_count)
+    baz = wrapped_module.baz  # pylint: disable=unused-variable
+    self.assertEqual(2, mock_warning.call_count)
+    baz = wrapped_module.baz
+    self.assertEqual(2, mock_warning.call_count)
+
+    # Check that values stayed the same
+    self.assertEqual(module.foo, foo)
+    self.assertEqual(module.bar, bar)
+
+
+class LazyLoadingWrapperTest(test.TestCase):
+
+  def testLazyLoad(self):
+    module = MockModule('test')
+    apis = {'cmd': ('', 'cmd'), 'ABCMeta': ('abc', 'ABCMeta')}
+    wrapped_module = module_wrapper.TFModuleWrapper(
+        module, 'test', public_apis=apis, deprecation=False)
+    import cmd as _cmd  # pylint: disable=g-import-not-at-top
+    from abc import ABCMeta as _ABCMeta  # pylint: disable=g-import-not-at-top, g-importing-member
+    self.assertFalse(wrapped_module._fastdict_key_in('cmd'))
+    self.assertEqual(wrapped_module.cmd, _cmd)
+    # Verify that the APIs are added to the cache of FastModuleType object
+    self.assertTrue(wrapped_module._fastdict_key_in('cmd'))
+    self.assertFalse(wrapped_module._fastdict_key_in('ABCMeta'))
+    self.assertEqual(wrapped_module.ABCMeta, _ABCMeta)
+    self.assertTrue(wrapped_module._fastdict_key_in('ABCMeta'))
+
+  def testLazyLoadLocalOverride(self):
+    # Test that we can override and add fields to the wrapped module.
+    module = MockModule('test')
+    apis = {'cmd': ('', 'cmd')}
+    wrapped_module = module_wrapper.TFModuleWrapper(
+        module, 'test', public_apis=apis, deprecation=False)
+    import cmd as _cmd  # pylint: disable=g-import-not-at-top
+    self.assertEqual(wrapped_module.cmd, _cmd)
+    setattr(wrapped_module, 'cmd', 1)
+    setattr(wrapped_module, 'cgi', 2)
+    self.assertEqual(wrapped_module.cmd, 1)  # override
+    # Verify that the values are also updated in the cache
+    # of the FastModuleType object
+    self.assertEqual(wrapped_module._fastdict_get('cmd'), 1)
+    self.assertEqual(wrapped_module.cgi, 2)  # add
+    self.assertEqual(wrapped_module._fastdict_get('cgi'), 2)
+
+  def testLazyLoadDict(self):
+    # Test that we can override and add fields to the wrapped module.
+    module = MockModule('test')
+    apis = {'cmd': ('', 'cmd')}
+    wrapped_module = module_wrapper.TFModuleWrapper(
+        module, 'test', public_apis=apis, deprecation=False)
+    import cmd as _cmd  # pylint: disable=g-import-not-at-top
+    # At first cmd key does not exist in __dict__
+    self.assertNotIn('cmd', wrapped_module.__dict__)
+    # After it is referred (lazyloaded), it gets added to __dict__
+    wrapped_module.cmd  # pylint: disable=pointless-statement
+    self.assertEqual(wrapped_module.__dict__['cmd'], _cmd)
+    # When we call setattr, it also gets added to __dict__
+    setattr(wrapped_module, 'cmd2', _cmd)
+    self.assertEqual(wrapped_module.__dict__['cmd2'], _cmd)
+
+  def testLazyLoadWildcardImport(self):
+    # Test that public APIs are in __all__.
+    module = MockModule('test')
+    module._should_not_be_public = 5
+    apis = {'cmd': ('', 'cmd')}
+    wrapped_module = module_wrapper.TFModuleWrapper(
+        module, 'test', public_apis=apis, deprecation=False)
+    setattr(wrapped_module, 'hello', 1)
+    self.assertIn('hello', wrapped_module.__all__)
+    self.assertIn('cmd', wrapped_module.__all__)
+    self.assertNotIn('_should_not_be_public', wrapped_module.__all__)
+
+  def testLazyLoadCorrectLiteModule(self):
+    # If set, always load lite module from public API list.
+    module = MockModule('test')
+    apis = {'lite': ('', 'cmd')}
+    module.lite = 5
+    import cmd as _cmd  # pylint: disable=g-import-not-at-top
+    wrapped_module = module_wrapper.TFModuleWrapper(
+        module, 'test', public_apis=apis, deprecation=False, has_lite=True)
+    self.assertEqual(wrapped_module.lite, _cmd)
+
+  def testInitCachesAttributes(self):
+    module = MockModule('test')
+    wrapped_module = module_wrapper.TFModuleWrapper(module, 'test')
+    self.assertTrue(wrapped_module._fastdict_key_in('_fastdict_key_in'))
+    self.assertTrue(wrapped_module._fastdict_key_in('_tfmw_module_name'))
+    self.assertTrue(wrapped_module._fastdict_key_in('__all__'))
+
+  def testCompatV1APIInstrumenting(self):
+    self.assertFalse(module_wrapper.TFModuleWrapper.compat_v1_usage_recorded)
+    apis = {'cosh': ('', 'cmd')}
+
+    mock_tf = MockModule('tensorflow')
+    mock_tf_wrapped = module_wrapper.TFModuleWrapper(
+        mock_tf, 'test', public_apis=apis)
+    mock_tf_wrapped.cosh  # pylint: disable=pointless-statement
+    self.assertFalse(module_wrapper.TFModuleWrapper.compat_v1_usage_recorded)
+
+    mock_tf_v1 = MockModule('tensorflow.compat.v1')
+    mock_tf_v1_wrapped = module_wrapper.TFModuleWrapper(
+        mock_tf_v1, 'test', public_apis=apis)
+    self.assertFalse(module_wrapper.TFModuleWrapper.compat_v1_usage_recorded)
+    mock_tf_v1_wrapped.cosh  # pylint: disable=pointless-statement
+    self.assertTrue(module_wrapper.TFModuleWrapper.compat_v1_usage_recorded)
+
+    # 'Reset' the status before testing against 'tensorflow.compat.v2.compat.v1'
+    module_wrapper.TFModuleWrapper.compat_v1_usage_recorded = False
+    mock_tf_v2_v1 = mock_tf_v1 = MockModule('tensorflow.compat.v2.compat.v1')
+    mock_tf_v2_v1_wrapped = module_wrapper.TFModuleWrapper(
+        mock_tf_v2_v1, 'test', public_apis=apis)
+    self.assertFalse(module_wrapper.TFModuleWrapper.compat_v1_usage_recorded)
+    mock_tf_v2_v1_wrapped.cosh  # pylint: disable=pointless-statement
+    self.assertTrue(module_wrapper.TFModuleWrapper.compat_v1_usage_recorded)
+
+  def testDelAttr(self):
+    module = MockModule('test')
+    wrapped_module = module_wrapper.TFModuleWrapper(module, 'test')
+    setattr(wrapped_module, 'foo', 1)
+    self.assertEqual(wrapped_module.foo, 1)
+    delattr(wrapped_module, 'foo')
+    self.assertFalse(hasattr(wrapped_module, 'foo'))
+    # Try setting the attr again
+    setattr(wrapped_module, 'foo', 1)
+    self.assertEqual(wrapped_module.foo, 1)
+    delattr(wrapped_module, 'foo')
+    self.assertFalse(hasattr(wrapped_module, 'foo'))
+
+
+class PickleTest(test.TestCase):
+
+  def testPickleSubmodule(self):
+    name = PickleTest.__module__  # The current module is a submodule.
+    module = module_wrapper.TFModuleWrapper(MockModule(name), name)
+    restored = pickle.loads(pickle.dumps(module))
+    self.assertEqual(restored.__name__, name)
+    self.assertIsNotNone(restored.PickleTest)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/nest.cc

@@ -0,0 +1,186 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/python/util/nest.h"
+
+#include <Python.h>
+
+#include <cstddef>
+#include <string>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+#include "tensorflow/python/util/util.h"
+
+namespace tensorflow {
+
+namespace {
+
+// Gets a string representation of the input object.
+//
+// Args:
+//   o: a python object.
+//   length: If set to negative, the whole string is returned. Otherwise, the
+//       string gets clipped to 'length' in size.
+//
+// Returns:
+//   A string representation.
+std::string PyObject_ToString(PyObject* o, int length = -1) {
+  auto str_o = make_safe(PyObject_Str(o));
+  if (!str_o) {
+    PyErr_Clear();
+    return "<failed to convert object to string>";
+  }
+  const char* utf8 = PyUnicode_AsUTF8(str_o.get());
+  if (utf8 == nullptr) {
+    PyErr_Clear();
+    return "<failed to convert unicode to utf8>";
+  }
+  std::string str(utf8);
+  if (length < 0 || str.size() <= length) {
+    return str;
+  }
+  absl::string_view str_piece(str);
+  return absl::StrCat(str_piece.substr(0, length), "...");
+}
+
+// Gets a list of keys from a dict or mapping type object.
+//
+// Args:
+//   o: a dictionary or mapping type object.
+//
+// Returns:
+//   A new reference to a list.
+//
+// Raises:
+//   TypeError: if `o` is not a dict or mapping type object.
+PyObject* GetKeysFromDictOrMapping(PyObject* o) {
+  if (PyDict_Check(o)) {
+    return PyDict_Keys(o);
+  } else if (PyMapping_Check(o)) {
+    return PyMapping_Keys(o);
+  } else {
+    auto* o_type = Py_TYPE(o);
+    PyErr_SetString(
+        PyExc_TypeError,
+        absl::StrCat("Expecting a type compatible with dict or mapping, got '",
+                     o_type->tp_name, "'")
+            .c_str());
+    return nullptr;
+  }
+}
+
+}  // namespace
+
+PyObject* FlattenDictItems(PyObject* dict) {
+  if (!PyDict_Check(dict) && !swig::IsMapping(dict)) {
+    PyErr_SetString(
+        PyExc_TypeError,
+        absl::StrCat("FlattenDictItems: 'dict' must be a dictionary or ",
+                     "collection.Mapping type object, instead of '",
+                     Py_TYPE(dict)->tp_name, "'.")
+            .c_str());
+    return nullptr;
+  }
+  PyObject* flat_dictionary = PyDict_New();
+  if (flat_dictionary == nullptr) {
+    return nullptr;
+  }
+  auto keys = make_safe(GetKeysFromDictOrMapping(dict));
+  if (!keys) {
+    Py_DecRef(flat_dictionary);
+    return nullptr;
+  }
+  for (size_t i = 0; i < PyList_Size(keys.get()); ++i) {
+    auto* key = PyList_GetItem(keys.get(), i);
+    if (key == nullptr) {
+      Py_DecRef(flat_dictionary);
+      return nullptr;
+    }
+    // We use a general approach in case 'dict' is a PyMapping type,
+    // but not a PyDict type.
+    auto value = make_safe(PyObject_GetItem(dict, key));
+    if (!value) {
+      Py_DecRef(flat_dictionary);
+      return nullptr;
+    }
+    if (swig::IsNested(key)) {
+      // The dict might contain list - list pairs.
+      auto flat_keys = make_safe(swig::Flatten(key, false));
+      if (!flat_keys) {
+        Py_DecRef(flat_dictionary);
+        return nullptr;
+      }
+      auto flat_values = make_safe(swig::Flatten(value.get(), false));
+      if (!flat_values) {
+        Py_DecRef(flat_dictionary);
+        return nullptr;
+      }
+      size_t flat_keys_sz = PyList_Size(flat_keys.get());
+      size_t flat_values_sz = PyList_Size(flat_values.get());
+      if (flat_keys_sz != flat_values_sz) {
+        PyErr_SetString(
+            PyExc_ValueError,
+            tensorflow::strings::StrCat(
+                "Could not flatten dictionary. Key had ", flat_keys_sz,
+                " elements, but value had ", flat_values_sz,
+                " elements. Key: ", PyObject_ToString(flat_keys.get()),
+                ", value: ", PyObject_ToString(flat_values.get()), ".")
+                .c_str());
+        Py_DecRef(flat_dictionary);
+        return nullptr;
+      }
+      for (size_t i = 0; i < flat_keys_sz; ++i) {
+        auto* flat_key = PyList_GetItem(flat_keys.get(), i);
+        if (flat_key == nullptr) {
+          Py_DecRef(flat_dictionary);
+          return nullptr;
+        }
+        auto* flat_value = PyList_GetItem(flat_values.get(), i);
+        if (flat_value == nullptr) {
+          Py_DecRef(flat_dictionary);
+          return nullptr;
+        }
+        if (PyDict_GetItem(flat_dictionary, flat_key) != nullptr) {
+          PyErr_SetString(
+              PyExc_ValueError,
+              absl::StrCat(
+                  "Cannot flatten dict because this key is not unique: ",
+                  PyObject_ToString(flat_key))
+                  .c_str());
+          Py_DecRef(flat_dictionary);
+          return nullptr;
+        }
+        PyDict_SetItem(flat_dictionary, flat_key, flat_value);
+      }
+    } else {
+      if (PyDict_GetItem(flat_dictionary, key) != nullptr) {
+        PyErr_SetString(
+            PyExc_ValueError,
+            absl::StrCat("Cannot flatten dict because this key is not unique: ",
+                         PyObject_ToString(key))
+                .c_str());
+        Py_DecRef(flat_dictionary);
+        return nullptr;
+      }
+      PyDict_SetItem(flat_dictionary, key, value.get());
+    }
+  }
+  return flat_dictionary;
+}
+
+}  // namespace tensorflow

tensorflow/python/util/nest.h

@@ -0,0 +1,37 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_PYTHON_UTIL_NEST_H_
+#define TENSORFLOW_PYTHON_UTIL_NEST_H_
+
+#include <Python.h>
+
+namespace tensorflow {
+// Returns a dictionary with flattened keys and values.
+//
+// Args:
+//   dict: the dictionary to zip
+//
+// Returns:
+//   An new reference to the zipped dictionary.
+//
+// Raises:
+//   TypeError: If the input is not a dictionary.
+//   ValueError: If any key and value do not have the same structure layout, or
+//       if keys are not unique.
+PyObject* FlattenDictItems(PyObject* dict);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_UTIL_NEST_H_

tensorflow/python/util/nest_wrapper.cc

@@ -0,0 +1,35 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/util/nest.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_nest, m) {
+  m.doc() = R"pbdoc(
+    _pywrap_nest
+    -----
+  )pbdoc";
+  m.def(
+      "FlattenDictItems",
+      [](const py::handle& dict) {
+        return tensorflow::PyoOrThrow(tensorflow::FlattenDictItems(dict.ptr()));
+      },
+      R"pbdoc(
+    Returns a dictionary with flattened keys and values.
+  )pbdoc");
+}

tensorflow/python/util/numpy_compat.py

@@ -0,0 +1,143 @@
+# 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.
+# ==============================================================================
+
+"""Functions for NumPy 1.x vs. 2.x compatibility."""
+
+import numpy as np
+
+
+def np_array(values, dtype=None, copy=True, order='K'):
+  """Creates a NumPy array containing input values.
+
+  It will make a copy of the object.
+
+  In NumPy 2.x and later, strict type casting can lead to errors when values
+  overflow the specified dtype. This function addresses this by replacing direct
+  np.array(..., dtype=...) calls with np.array(...).astype(...). This allows for
+  intended overflows, aligning with the behavior of older NumPy versions.
+
+  Args:
+    values: Array_like objects. E.g., a python list, tuple, or an object whose
+      __array__ method returns an array.
+    dtype: The desired numpy data type for the array.
+    copy: Bool. If True (default), then the array data is copied. If None, a
+      copy will only be made if __array__ returns a copy, if obj is a nested
+      sequence, or if a copy is needed to satisfy any of the other requirements
+      (dtype, order, etc.). Note that any copy of the data is shallow, i.e., for
+      arrays with object dtype, the new array will point to the same objects.
+      For False it raises a ValueError if a copy cannot be avoided.
+    order: {‘K’, ‘A’, ‘C’, ‘F’}.
+
+  Returns:
+    A NumPy array with the specified data type.
+  """
+  if dtype is not None and np.issubdtype(dtype, np.number):
+    return np.array(values, copy=copy, order=order).astype(dtype)
+  else:
+    return np.array(values, dtype=dtype, copy=copy, order=order)
+
+
+def np_asarray(values, dtype=None, order=None, copy=None):
+  """Converts input values to a NumPy array.
+
+  It will not make a copy.
+
+  In NumPy 2.x and later, strict type casting can lead to errors when values
+  overflow the specified dtype. This function addresses this by replacing direct
+  np.array(..., dtype=...) calls with np.array(...).astype(...). This allows for
+  intended overflows, aligning with the behavior of older NumPy versions.
+
+  Args:
+    values: Array_like objects. E.g., a python list, tuple, or an object whose
+      __array__ method returns an array.
+    dtype: The desired numpy data type for the array.
+    order: {‘C’, ‘F’, ‘A’, ‘K’}.
+    copy: bool. If True, then the object is copied. If None then the object is
+      copied only if needed, i.e. if __array__ returns a copy, if obj is a
+      nested sequence, or if a copy is needed to satisfy any of the other
+      requirements (dtype, order, etc.). For False it raises a ValueError if a
+      copy cannot be avoided.
+
+  Returns:
+    A NumPy array with the specified data type.
+  """
+  if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
+    if dtype is not None and np.issubdtype(dtype, np.number):
+      return np.asarray(values, order=order, copy=copy).astype(dtype, copy=copy)
+    else:
+      return np.asarray(values, dtype=dtype, order=order, copy=copy)
+  else:
+    return np.asarray(values, dtype=dtype, order=order)
+
+
+def np_where(condition, x=None, y=None):
+  """Return elements chosen from x or y depending on condition.
+
+  When only condition is provided, np.where(condition) is a shorthand for
+  np.asarray(condition).nonzero(). See
+  https://numpy.org/doc/stable/reference/generated/numpy.where.html. NumPy
+  2.1.0rc0 disallows 0D input arrays in nonzero, so np.atleast_1d is used here
+  to remain compatible with NumPy 1.x. See
+  https://github.com/numpy/numpy/pull/26268.
+
+  Args:
+    condition: Array_like, bool. Where True, yield x, otherwise yield y.
+    x: Array_like. Values from which to choose. x, y and condition need to be
+    broadcastable to some shape.
+    y: Array_like. Values from which to choose. x, y and condition need to be
+    broadcastable to some shape.
+
+  Returns:
+    An array with elements from x where condition is True, and elements from y
+    elsewhere. Or the indices of the elements that are non-zero.
+  """
+  if x is None and y is None:
+    if np.lib.NumpyVersion(np.__version__) >= '2.1.0.rc0':
+      return np.atleast_1d(np.asarray(condition)).nonzero()
+    return np.where(condition)
+  return np.where(condition, x, y)
+
+
+def np_reshape(a, /, shape=None, *, newshape=None, order='C', copy=None):
+  """Reshapes an array without changing its data.
+
+  NumPy 2.1.0rc1 added shape and copy arguments to numpy.reshape. See
+  https://github.com/numpy/numpy/pull/26292. Both newshape and shape keywords
+  are supported, but newshape is going to be deprecated. Use `shape` instead.
+
+  Besides, shape cannot be None now. See
+  https://github.com/numpy/numpy/blob/v2.1.0rc1/numpy/_core/fromnumeric.py#L309.
+  Previously, np.reshape with newshape=None returned a copy. To maintain this
+  behavior, we now use asarray to create an ndarray.
+
+  Args:
+    a: Array_like. Array to be reshaped.
+    shape: The new shape of the array.
+    newshape: The new shape of the array (deprecated).
+    order: {‘C’, ‘F’, ‘K’}.
+    copy: bool. If True, then the array data is copied. If None, a copy will
+    only be made if it’s required by order. For False it raises a ValueError if
+    a copy cannot be avoided.
+
+  Returns:
+    This will be a new view object if possible; otherwise, it will be a copy.
+  """
+  if shape is None:
+    shape = newshape
+  if np.lib.NumpyVersion(np.__version__) >= '2.1.0.rc0':
+    if shape is None and newshape is None:
+      return np.asarray(a, order=order, copy=copy)
+    return np.reshape(a, shape, order=order, copy=copy)
+  return np.reshape(a, shape, order=order)

tensorflow/python/util/numpy_compat_test.py

@@ -0,0 +1,61 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""numpy_compat tests."""
+
+
+import numpy as np
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import numpy_compat
+
+
+class NumpyCompatCopyBehaviorTest(test.TestCase):
+
+  def test_no_copy_new_vs_old(self):
+
+    # Define old_np_asarray to replicate the old code that used .astype(dtype)
+    # WITHOUT passing `copy=copy`.
+    def old_np_asarray(values, dtype=None, order=None, copy=None):
+      if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
+        if dtype is not None and np.issubdtype(dtype, np.number):
+          return np.asarray(values, order=order, copy=copy).astype(dtype)
+        else:
+          return np.asarray(values, dtype=dtype, order=order, copy=copy)
+      else:
+        return np.asarray(values, dtype=dtype, order=order)
+
+    # Test array
+    x = np.array([1, 2, 3], dtype=np.float32)
+
+    # Expect old numpy 2.x code to always copy even when copy=None
+    y_old = old_np_asarray(x, dtype=np.float32, copy=None)
+    if np.lib.NumpyVersion(np.__version__) >= '2.0.0.dev0':
+      self.assertIsNot(
+          y_old,
+          x,
+          msg='Old code did NOT copy, but we expect it to always copy.',
+      )
+
+    # Expect new code to reuse the array if copy=None
+    y_new = numpy_compat.np_asarray(x, dtype=np.float32, copy=None)
+    self.assertIs(
+        y_new,
+        x,
+        msg='New code did copy, but we expect it NOT to copy since copy=None.',
+    )
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/object_identity.py

@@ -0,0 +1,265 @@
+"""Utilities for collecting objects based on "is" comparison."""
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Any, Set
+import weakref
+
+from tensorflow.python.util.compat import collections_abc
+
+
+# LINT.IfChange
+class _ObjectIdentityWrapper:
+  """Wraps an object, mapping __eq__ on wrapper to "is" on wrapped.
+
+  Since __eq__ is based on object identity, it's safe to also define __hash__
+  based on object ids. This lets us add unhashable types like trackable
+  _ListWrapper objects to object-identity collections.
+  """
+
+  __slots__ = ["_wrapped", "__weakref__"]
+
+  def __init__(self, wrapped):
+    self._wrapped = wrapped
+
+  @property
+  def unwrapped(self):
+    return self._wrapped
+
+  def _assert_type(self, other):
+    if not isinstance(other, _ObjectIdentityWrapper):
+      raise TypeError("Cannot compare wrapped object with unwrapped object")
+
+  def __lt__(self, other):
+    self._assert_type(other)
+    return id(self._wrapped) < id(other._wrapped)  # pylint: disable=protected-access
+
+  def __gt__(self, other):
+    self._assert_type(other)
+    return id(self._wrapped) > id(other._wrapped)  # pylint: disable=protected-access
+
+  def __eq__(self, other):
+    if other is None:
+      return False
+    self._assert_type(other)
+    return self._wrapped is other._wrapped  # pylint: disable=protected-access
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
+  def __hash__(self):
+    # Wrapper id() is also fine for weakrefs. In fact, we rely on
+    # id(weakref.ref(a)) == id(weakref.ref(a)) and weakref.ref(a) is
+    # weakref.ref(a) in _WeakObjectIdentityWrapper.
+    return id(self._wrapped)
+
+  def __repr__(self):
+    return "<{} wrapping {!r}>".format(type(self).__name__, self._wrapped)
+
+
+class _WeakObjectIdentityWrapper(_ObjectIdentityWrapper):
+
+  __slots__ = ()
+
+  def __init__(self, wrapped):
+    super(_WeakObjectIdentityWrapper, self).__init__(weakref.ref(wrapped))
+
+  @property
+  def unwrapped(self):
+    return self._wrapped()
+
+
+class Reference(_ObjectIdentityWrapper):
+  """Reference that refers an object.
+
+  ```python
+  x = [1]
+  y = [1]
+
+  x_ref1 = Reference(x)
+  x_ref2 = Reference(x)
+  y_ref2 = Reference(y)
+
+  print(x_ref1 == x_ref2)
+  ==> True
+
+  print(x_ref1 == y)
+  ==> False
+  ```
+  """
+
+  __slots__ = ()
+
+  # Disabling super class' unwrapped field.
+  unwrapped = property()
+
+  def deref(self):
+    """Returns the referenced object.
+
+    ```python
+    x_ref = Reference(x)
+    print(x is x_ref.deref())
+    ==> True
+    ```
+    """
+    return self._wrapped
+
+
+class ObjectIdentityDictionary(collections_abc.MutableMapping):
+  """A mutable mapping data structure which compares using "is".
+
+  This is necessary because we have trackable objects (_ListWrapper) which
+  have behavior identical to built-in Python lists (including being unhashable
+  and comparing based on the equality of their contents by default).
+  """
+
+  __slots__ = ["_storage"]
+
+  def __init__(self):
+    self._storage = {}
+
+  def _wrap_key(self, key):
+    return _ObjectIdentityWrapper(key)
+
+  def __getitem__(self, key):
+    return self._storage[self._wrap_key(key)]
+
+  def __setitem__(self, key, value):
+    self._storage[self._wrap_key(key)] = value
+
+  def __delitem__(self, key):
+    del self._storage[self._wrap_key(key)]
+
+  def __len__(self):
+    return len(self._storage)
+
+  def __iter__(self):
+    for key in self._storage:
+      yield key.unwrapped
+
+  def __repr__(self):
+    return "ObjectIdentityDictionary(%s)" % repr(self._storage)
+
+
+class ObjectIdentityWeakKeyDictionary(ObjectIdentityDictionary):
+  """Like weakref.WeakKeyDictionary, but compares objects with "is"."""
+
+  __slots__ = ["__weakref__"]
+
+  def _wrap_key(self, key):
+    return _WeakObjectIdentityWrapper(key)
+
+  def __len__(self):
+    # Iterate, discarding old weak refs
+    return len(list(self._storage))
+
+  def __iter__(self):
+    keys = self._storage.keys()
+    for key in keys:
+      unwrapped = key.unwrapped
+      if unwrapped is None:
+        del self[key]
+      else:
+        yield unwrapped
+
+
+class ObjectIdentitySet(collections_abc.MutableSet):
+  """Like the built-in set, but compares objects with "is"."""
+
+  __slots__ = ["_storage", "__weakref__"]
+
+  def __init__(self, *args):
+    self._storage = set(self._wrap_key(obj) for obj in list(*args))
+
+  def __le__(self, other: Set[Any]) -> bool:
+    if not isinstance(other, Set):
+      return NotImplemented
+    if len(self) > len(other):
+      return False
+    for item in self._storage:
+      if item not in other:
+        return False
+    return True
+
+  def __ge__(self, other: Set[Any]) -> bool:
+    if not isinstance(other, Set):
+      return NotImplemented
+    if len(self) < len(other):
+      return False
+    for item in other:
+      if item not in self:
+        return False
+    return True
+
+  @staticmethod
+  def _from_storage(storage):
+    result = ObjectIdentitySet()
+    result._storage = storage  # pylint: disable=protected-access
+    return result
+
+  def _wrap_key(self, key):
+    return _ObjectIdentityWrapper(key)
+
+  def __contains__(self, key):
+    return self._wrap_key(key) in self._storage
+
+  def discard(self, key):
+    self._storage.discard(self._wrap_key(key))
+
+  def add(self, key):
+    self._storage.add(self._wrap_key(key))
+
+  def update(self, items):
+    self._storage.update([self._wrap_key(item) for item in items])
+
+  def clear(self):
+    self._storage.clear()
+
+  def intersection(self, items):
+    return self._storage.intersection([self._wrap_key(item) for item in items])
+
+  def difference(self, items):
+    return ObjectIdentitySet._from_storage(
+        self._storage.difference([self._wrap_key(item) for item in items]))
+
+  def __len__(self):
+    return len(self._storage)
+
+  def __iter__(self):
+    keys = list(self._storage)
+    for key in keys:
+      yield key.unwrapped
+
+
+class ObjectIdentityWeakSet(ObjectIdentitySet):
+  """Like weakref.WeakSet, but compares objects with "is"."""
+
+  __slots__ = ()
+
+  def _wrap_key(self, key):
+    return _WeakObjectIdentityWrapper(key)
+
+  def __len__(self):
+    # Iterate, discarding old weak refs
+    return len([_ for _ in self])
+
+  def __iter__(self):
+    keys = list(self._storage)
+    for key in keys:
+      unwrapped = key.unwrapped
+      if unwrapped is None:
+        self.discard(key)
+      else:
+        yield unwrapped
+# LINT.ThenChange(//tensorflow/python/keras/utils/object_identity.py)

tensorflow/python/util/object_identity_test.py

@@ -0,0 +1,101 @@
+# 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.
+# ==============================================================================
+"""Unit tests for object_identity."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import nest
+from tensorflow.python.util import object_identity
+
+
+class ObjectIdentityWrapperTest(test.TestCase):
+
+  def testWrapperNotEqualToWrapped(self):
+    class SettableHash(object):
+
+      def __init__(self):
+        self.hash_value = 8675309
+
+      def __hash__(self):
+        return self.hash_value
+
+    o = SettableHash()
+    wrap1 = object_identity._ObjectIdentityWrapper(o)
+    wrap2 = object_identity._ObjectIdentityWrapper(o)
+
+    self.assertEqual(wrap1, wrap1)
+    self.assertEqual(wrap1, wrap2)
+    self.assertEqual(o, wrap1.unwrapped)
+    self.assertEqual(o, wrap2.unwrapped)
+    with self.assertRaises(TypeError):
+      bool(o == wrap1)
+    with self.assertRaises(TypeError):
+      bool(wrap1 != o)
+
+    self.assertNotIn(o, set([wrap1]))
+    o.hash_value = id(o)
+    # Since there is now a hash collision we raise an exception
+    with self.assertRaises(TypeError):
+      bool(o in set([wrap1]))
+
+  def testNestFlatten(self):
+    a = object_identity._ObjectIdentityWrapper('a')
+    b = object_identity._ObjectIdentityWrapper('b')
+    c = object_identity._ObjectIdentityWrapper('c')
+    flat = nest.flatten([[[(a, b)]], c])
+    self.assertEqual(flat, [a, b, c])
+
+  def testNestMapStructure(self):
+    k = object_identity._ObjectIdentityWrapper('k')
+    v1 = object_identity._ObjectIdentityWrapper('v1')
+    v2 = object_identity._ObjectIdentityWrapper('v2')
+    struct = nest.map_structure(lambda a, b: (a, b), {k: v1}, {k: v2})
+    self.assertEqual(struct, {k: (v1, v2)})
+
+
+class ObjectIdentitySetTest(test.TestCase):
+
+  def testDifference(self):
+
+    class Element(object):
+      pass
+
+    a = Element()
+    b = Element()
+    c = Element()
+    set1 = object_identity.ObjectIdentitySet([a, b])
+    set2 = object_identity.ObjectIdentitySet([b, c])
+    diff_set = set1.difference(set2)
+    self.assertIn(a, diff_set)
+    self.assertNotIn(b, diff_set)
+    self.assertNotIn(c, diff_set)
+
+  def testDiscard(self):
+    a = object()
+    b = object()
+    set1 = object_identity.ObjectIdentitySet([a, b])
+    set1.discard(a)
+    self.assertIn(b, set1)
+    self.assertNotIn(a, set1)
+
+  def testClear(self):
+    a = object()
+    b = object()
+    set1 = object_identity.ObjectIdentitySet([a, b])
+    set1.clear()
+    self.assertLen(set1, 0)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/port_wrapper.cc

@@ -0,0 +1,33 @@
+/* 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/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/core/util/port.h"
+
+PYBIND11_MODULE(_pywrap_util_port, m) {
+  m.def("IsGoogleCudaEnabled", tensorflow::IsGoogleCudaEnabled);
+  m.def("IsBuiltWithROCm", tensorflow::IsBuiltWithROCm);
+  m.def("IsBuiltWithXLA", tensorflow::IsBuiltWithXLA);
+  m.def("IsBuiltWithNvcc", tensorflow::IsBuiltWithNvcc);
+  m.def("IsAArch32Available", tensorflow::IsAArch32Available);
+  m.def("IsAArch64Available", tensorflow::IsAArch64Available);
+  m.def("IsPowerPCAvailable", tensorflow::IsPowerPCAvailable);
+  m.def("IsSystemZAvailable", tensorflow::IsSystemZAvailable);
+  m.def("IsX86Available", tensorflow::IsX86Available);
+  m.def("GpuSupportsHalfMatMulAndConv",
+        tensorflow::GpuSupportsHalfMatMulAndConv);
+  m.def("IsMklEnabled", tensorflow::IsMklEnabled);
+}

tensorflow/python/util/protobuf/BUILD

@@ -0,0 +1,86 @@
+# Tensorflow protobuf utility package
+
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "tf_py_strict_test")
+load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")  # @unused
+
+visibility = [
+    "//third_party/cloud_tpu/convergence_tools:__subpackages__",
+    "//tensorflow:internal",
+    "//tensorflow/lite/toco/python:__pkg__",
+    "//tensorflow_models:__subpackages__",
+    "//tensorflow_model_optimization:__subpackages__",
+    "//third_party/py/cleverhans:__subpackages__",
+    "//third_party/py/launchpad:__subpackages__",
+    "//third_party/py/reverb:__subpackages__",
+    "//third_party/py/neural_structured_learning:__subpackages__",
+    "//third_party/py/tensorflow_examples:__subpackages__",
+    "//third_party/py/tf_agents:__subpackages__",  # For benchmarks.
+    "//third_party/py/tf_slim:__subpackages__",
+    "//third_party/py/tensorflow_docs:__subpackages__",
+]
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = visibility,
+    licenses = ["notice"],
+)
+
+tf_proto_library(
+    name = "compare_test_proto",
+    testonly = 1,
+    srcs = ["compare_test.proto"],
+)
+
+tf_py_strict_test(
+    name = "protobuf_compare_test",
+    size = "small",
+    srcs = ["compare_test.py"],
+    main = "compare_test.py",
+    tags = ["no_pip"],  # compare_test_pb2 proto is not available in pip.
+    deps = [
+        ":compare_test_proto_py",
+        ":protobuf",
+        "//tensorflow/python/platform:test",
+    ],
+)
+
+filegroup(
+    name = "compare_test_proto_src",
+    srcs = ["compare_test.proto"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "compare_test_py_pb2",
+#     testonly = 1,
+#     has_services = 0,
+#     deps = [":compare_test_proto"],
+# )
+# copybara:uncomment_end
+
+py_library(
+    name = "protobuf",
+    srcs = glob(
+        ["*.py"],
+        exclude = ["*_test.py"],
+    ),
+    compatible_with = get_compatible_with_portable(),
+    strict_deps = True,
+    visibility = visibility + [
+        "//tensorflow:__pkg__",
+        "//third_party/py/tensorflow_core:__subpackages__",
+        "//third_party/py/tensorflow_data_validation:__subpackages__",
+        "//third_party/py/tf_agents:__subpackages__",
+        "//third_party/py/tfx:__subpackages__",
+    ],
+    deps = [
+        # global_test_configuration is added here because all major tests depend on this
+        # library. It isn't possible to add these test dependencies via tensorflow.bzl's
+        # py_test because not all tensorflow tests use tensorflow.bzl's py_test.
+        "//tensorflow/python:global_test_configuration",
+        "@com_google_protobuf//:protobuf_python",
+        "//tensorflow/python/util:compat",
+    ],
+)

tensorflow/python/util/protobuf/compare.py

@@ -0,0 +1,377 @@
+# 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 comparing proto2 messages in Python.
+
+ProtoEq() compares two proto2 messages for equality.
+
+ClearDefaultValuedFields() recursively clears the fields that are set to their
+default values. This is useful for comparing protocol buffers where the
+semantics of unset fields and default valued fields are the same.
+
+assertProtoEqual() is useful for unit tests.  It produces much more helpful
+output than assertEqual() for proto2 messages, e.g. this:
+
+  outer {
+    inner {
+-     strings: "x"
+?               ^
++     strings: "y"
+?               ^
+    }
+  }
+
+...compared to the default output from assertEqual() that looks like this:
+
+AssertionError: <my.Msg object at 0x9fb353c> != <my.Msg object at 0x9fb35cc>
+
+Call it inside your unit test's googletest.TestCase subclasses like this:
+
+  from tensorflow.python.util.protobuf import compare
+
+  class MyTest(googletest.TestCase):
+    ...
+    def testXXX(self):
+      ...
+      compare.assertProtoEqual(self, a, b)
+
+Alternatively:
+
+  from tensorflow.python.util.protobuf import compare
+
+  class MyTest(compare.ProtoAssertions, googletest.TestCase):
+    ...
+    def testXXX(self):
+      ...
+      self.assertProtoEqual(a, b)
+"""
+
+import collections.abc as collections_abc
+import difflib
+import math
+
+from google.protobuf import descriptor
+from google.protobuf import descriptor_pool
+from google.protobuf import message
+from google.protobuf import text_format
+
+
+# TODO(alankelly): Distinguish between signalling and quiet NaNs.
+def isClose(x, y, relative_tolerance):  # pylint: disable=invalid-name
+  """Returns True if x is close to y given the relative tolerance or if x and y are both inf, both -inf, or both NaNs.
+
+  This function does not distinguish between signalling and non-signalling NaN.
+
+  Args:
+    x: float value to be compared
+    y: float value to be compared
+    relative_tolerance: float. The allowable difference between the two values
+      being compared is determined by multiplying the relative tolerance by the
+      maximum of the two values. If this is not provided, then all floats are
+      compared using string comparison.
+  """
+  # NaNs are considered equal.
+  if math.isnan(x) or math.isnan(y):
+    return math.isnan(x) == math.isnan(y)
+
+  if math.isinf(x) or math.isinf(y):
+    return x == y
+
+  return abs(x - y) <= relative_tolerance * max(abs(x), abs(y))
+
+
+def checkFloatEqAndReplace(self, expected, actual, relative_tolerance):  # pylint: disable=invalid-name
+  """Recursively replaces the floats in actual with those in expected iff they are approximately equal.
+
+  This is done because string equality will consider values such as 5.0999999999
+  and 5.1 as not being equal, despite being extremely close.
+
+  Args:
+    self: googletest.TestCase
+    expected: expected values
+    actual: actual values
+    relative_tolerance: float, relative tolerance.
+  """
+
+  for expected_fields, actual_fields in zip(
+      expected.ListFields(), actual.ListFields()
+  ):
+    is_repeated = True
+    expected_desc, expected_values = expected_fields
+    actual_values = actual_fields[1]
+    if expected_desc.label != descriptor.FieldDescriptor.LABEL_REPEATED:
+      is_repeated = False
+      expected_values = [expected_values]
+      actual_values = [actual_values]
+
+    if (
+        expected_desc.type == descriptor.FieldDescriptor.TYPE_FLOAT
+        or expected_desc.type == descriptor.FieldDescriptor.TYPE_DOUBLE
+    ):
+      for i, (x, y) in enumerate(zip(expected_values, actual_values)):
+        # Replace the actual value with the expected value if the test passes,
+        # otherwise leave it and let it fail in the next test so that the error
+        # message is nicely formatted
+        if isClose(x, y, relative_tolerance):
+          if is_repeated:
+            getattr(actual, actual_fields[0].name)[i] = x
+          else:
+            setattr(actual, actual_fields[0].name, x)
+
+    if (
+        expected_desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE
+        or expected_desc.type == descriptor.FieldDescriptor.TYPE_GROUP
+    ):
+      if (
+          expected_desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE
+          and expected_desc.message_type.has_options
+          and expected_desc.message_type.GetOptions().map_entry
+      ):
+        # This is a map, only recurse if it has type message type.
+        if (
+            expected_desc.message_type.fields_by_number[2].type
+            == descriptor.FieldDescriptor.TYPE_MESSAGE
+        ):
+          for e_v, a_v in zip(
+              iter(expected_values.values()), iter(actual_values.values())
+          ):
+            checkFloatEqAndReplace(
+                self,
+                expected=e_v,
+                actual=a_v,
+                relative_tolerance=relative_tolerance,
+            )
+      else:
+        for v, a in zip(expected_values, actual_values):
+          # recursive step
+          checkFloatEqAndReplace(
+              self, expected=v, actual=a, relative_tolerance=relative_tolerance
+          )
+
+
+def assertProtoEqual(
+    self,
+    a,
+    b,
+    check_initialized=True,
+    normalize_numbers=False,
+    msg=None,
+    relative_tolerance=None,
+):  # pylint: disable=invalid-name(
+  """Fails with a useful error if a and b aren't equal.
+
+  Comparison of repeated fields matches the semantics of
+  unittest.TestCase.assertEqual(), ie order and extra duplicates fields matter.
+
+  Args:
+    self: googletest.TestCase
+    a: proto2 PB instance, or text string representing one.
+    b: proto2 PB instance -- message.Message or subclass thereof.
+    check_initialized: boolean, whether to fail if either a or b isn't
+      initialized.
+    normalize_numbers: boolean, whether to normalize types and precision of
+      numbers before comparison.
+    msg: if specified, is used as the error message on failure.
+    relative_tolerance: float, relative tolerance. If this is not provided, then
+      all floats are compared using string comparison otherwise, floating point
+      comparisons are done using the relative tolerance provided.
+  """
+  pool = descriptor_pool.Default()
+  if isinstance(a, str):
+    a = text_format.Parse(a, b.__class__(), descriptor_pool=pool)
+
+  for pb in a, b:
+    if check_initialized:
+      errors = pb.FindInitializationErrors()
+      if errors:
+        self.fail('Initialization errors: %s\n%s' % (errors, pb))
+    if normalize_numbers:
+      NormalizeNumberFields(pb)
+
+  if relative_tolerance is not None:
+    checkFloatEqAndReplace(
+        self, expected=b, actual=a, relative_tolerance=relative_tolerance
+    )
+
+  a_str = text_format.MessageToString(a, descriptor_pool=pool)
+  b_str = text_format.MessageToString(b, descriptor_pool=pool)
+
+  # Some Python versions would perform regular diff instead of multi-line
+  # diff if string is longer than 2**16. We substitute this behavior
+  # with a call to unified_diff instead to have easier-to-read diffs.
+  # For context, see: https://bugs.python.org/issue11763.
+  if len(a_str) < 2**16 and len(b_str) < 2**16:
+    self.assertMultiLineEqual(a_str, b_str, msg=msg)
+  else:
+    diff = ''.join(
+        difflib.unified_diff(a_str.splitlines(True), b_str.splitlines(True)))
+    if diff:
+      self.fail('%s :\n%s' % (msg, diff))
+
+
+def NormalizeNumberFields(pb):
+  """Normalizes types and precisions of number fields in a protocol buffer.
+
+  Due to subtleties in the python protocol buffer implementation, it is possible
+  for values to have different types and precision depending on whether they
+  were set and retrieved directly or deserialized from a protobuf. This function
+  normalizes integer values to ints and longs based on width, 32-bit floats to
+  five digits of precision to account for python always storing them as 64-bit,
+  and ensures doubles are floating point for when they're set to integers.
+
+  Modifies pb in place. Recurses into nested objects.
+
+  Args:
+    pb: proto2 message.
+
+  Returns:
+    the given pb, modified in place.
+  """
+  for desc, values in pb.ListFields():
+    is_repeated = True
+    if desc.label != descriptor.FieldDescriptor.LABEL_REPEATED:
+      is_repeated = False
+      values = [values]
+
+    normalized_values = None
+
+    # We force 32-bit values to int and 64-bit values to long to make
+    # alternate implementations where the distinction is more significant
+    # (e.g. the C++ implementation) simpler.
+    if desc.type in (descriptor.FieldDescriptor.TYPE_INT64,
+                     descriptor.FieldDescriptor.TYPE_UINT64,
+                     descriptor.FieldDescriptor.TYPE_SINT64):
+      normalized_values = [int(x) for x in values]
+    elif desc.type in (descriptor.FieldDescriptor.TYPE_INT32,
+                       descriptor.FieldDescriptor.TYPE_UINT32,
+                       descriptor.FieldDescriptor.TYPE_SINT32,
+                       descriptor.FieldDescriptor.TYPE_ENUM):
+      normalized_values = [int(x) for x in values]
+    elif desc.type == descriptor.FieldDescriptor.TYPE_FLOAT:
+      normalized_values = [round(x, 6) for x in values]
+    elif desc.type == descriptor.FieldDescriptor.TYPE_DOUBLE:
+      normalized_values = [round(float(x), 7) for x in values]
+
+    if normalized_values is not None:
+      if is_repeated:
+        pb.ClearField(desc.name)
+        getattr(pb, desc.name).extend(normalized_values)
+      else:
+        setattr(pb, desc.name, normalized_values[0])
+
+    if (desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE or
+        desc.type == descriptor.FieldDescriptor.TYPE_GROUP):
+      if (desc.type == descriptor.FieldDescriptor.TYPE_MESSAGE and
+          desc.message_type.has_options and
+          desc.message_type.GetOptions().map_entry):
+        # This is a map, only recurse if the values have a message type.
+        if (desc.message_type.fields_by_number[2].type ==
+            descriptor.FieldDescriptor.TYPE_MESSAGE):
+          for v in iter(values.values()):
+            NormalizeNumberFields(v)
+      else:
+        for v in values:
+          # recursive step
+          NormalizeNumberFields(v)
+
+  return pb
+
+
+def _IsMap(value):
+  return isinstance(value, collections_abc.Mapping)
+
+
+def _IsRepeatedContainer(value):
+  if isinstance(value, str):
+    return False
+  try:
+    iter(value)
+    return True
+  except TypeError:
+    return False
+
+
+def ProtoEq(a, b):
+  """Compares two proto2 objects for equality.
+
+  Recurses into nested messages. Uses list (not set) semantics for comparing
+  repeated fields, ie duplicates and order matter.
+
+  Args:
+    a: A proto2 message or a primitive.
+    b: A proto2 message or a primitive.
+
+  Returns:
+    `True` if the messages are equal.
+  """
+  def Format(pb):
+    """Returns a dictionary or unchanged pb bases on its type.
+
+    Specifically, this function returns a dictionary that maps tag
+    number (for messages) or element index (for repeated fields) to
+    value, or just pb unchanged if it's neither.
+
+    Args:
+      pb: A proto2 message or a primitive.
+    Returns:
+      A dict or unchanged pb.
+    """
+    if isinstance(pb, message.Message):
+      return dict((desc.number, value) for desc, value in pb.ListFields())
+    elif _IsMap(pb):
+      return dict(pb.items())
+    elif _IsRepeatedContainer(pb):
+      return dict(enumerate(list(pb)))
+    else:
+      return pb
+
+  a, b = Format(a), Format(b)
+
+  # Base case
+  if not isinstance(a, dict) or not isinstance(b, dict):
+    return a == b
+
+  # This list performs double duty: it compares two messages by tag value *or*
+  # two repeated fields by element, in order. the magic is in the format()
+  # function, which converts them both to the same easily comparable format.
+  for tag in sorted(set(a.keys()) | set(b.keys())):
+    if tag not in a or tag not in b:
+      return False
+    else:
+      # Recursive step
+      if not ProtoEq(a[tag], b[tag]):
+        return False
+
+  # Didn't find any values that differed, so they're equal!
+  return True
+
+
+class ProtoAssertions(object):
+  """Mix this into a googletest.TestCase class to get proto2 assertions.
+
+  Usage:
+
+  class SomeTestCase(compare.ProtoAssertions, googletest.TestCase):
+    ...
+    def testSomething(self):
+      ...
+      self.assertProtoEqual(a, b)
+
+  See module-level definitions for method documentation.
+  """
+
+  # pylint: disable=invalid-name
+  def assertProtoEqual(self, *args, **kwargs):
+    return assertProtoEqual(self, *args, **kwargs)

tensorflow/python/util/protobuf/compare_test.proto

@@ -0,0 +1,83 @@
+// 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.
+// ==============================================================================
+
+// Test messages used in compare_test.py.
+syntax = "proto2";
+
+package compare_test;
+
+option cc_enable_arenas = true;
+
+enum Enum {
+  A = 0;
+  B = 1;
+  C = 2;
+}
+
+message Small {
+  repeated string strings = 1;
+}
+
+message Medium {
+  repeated int32 int32s = 1;
+  repeated Small smalls = 2;
+  repeated group GroupA = 3 {
+    repeated group GroupB = 4 {
+      required string strings = 5;
+    }
+  }
+  repeated float floats = 6;
+}
+
+message Large {
+  optional string string_ = 1;
+  optional int64 int64_ = 2;
+  optional float float_ = 3;
+  optional bool bool_ = 4;
+  optional Enum enum_ = 5;
+  repeated int64 int64s = 6;
+  optional Medium medium = 7;
+  optional Small small = 8;
+  optional double double_ = 9;
+  optional WithMap with_map = 10;
+}
+
+message Labeled {
+  required int32 required = 1;
+  optional int32 optional = 2;
+}
+
+message WithMap {
+  map<int32, Small> value_message = 1;
+  map<string, string> value_string = 2;
+}
+
+message Floats {
+  optional float float_ = 1;
+  optional double double_ = 2;
+}
+
+message RepeatedFloats {
+  repeated float float_ = 1 [packed = true];
+  repeated double double_ = 2 [packed = true];
+}
+
+message NestedFloats {
+  optional Floats floats = 1;
+}
+
+message MapFloats {
+  map<int64, Floats> int_to_floats = 1;
+}

tensorflow/python/util/protobuf/compare_test.py

@@ -0,0 +1,533 @@
+# 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 python.util.protobuf.compare."""
+
+import copy
+import re
+import sys
+import textwrap
+
+
+from google.protobuf import text_format
+
+from tensorflow.python.platform import googletest
+from tensorflow.python.util.protobuf import compare
+from tensorflow.python.util.protobuf import compare_test_pb2
+
+
+def LargePbs(*args):
+  """Converts ASCII string Large PBs to messages."""
+  return [text_format.Merge(arg, compare_test_pb2.Large()) for arg in args]
+
+
+class ProtoEqTest(googletest.TestCase):
+
+  def assertNotEquals(self, a, b):
+    """Asserts that ProtoEq says a != b."""
+    a, b = LargePbs(a, b)
+    googletest.TestCase.assertEqual(self, compare.ProtoEq(a, b), False)
+
+  def assertEqual(self, a, b):
+    """Asserts that ProtoEq says a == b."""
+    a, b = LargePbs(a, b)
+    googletest.TestCase.assertEqual(self, compare.ProtoEq(a, b), True)
+
+  def testPrimitives(self):
+    googletest.TestCase.assertEqual(self, True, compare.ProtoEq('a', 'a'))
+    googletest.TestCase.assertEqual(self, False, compare.ProtoEq('b', 'a'))
+
+  def testEmpty(self):
+    self.assertEqual('', '')
+
+  def testPrimitiveFields(self):
+    self.assertNotEqual('string_: "a"', '')
+    self.assertEqual('string_: "a"', 'string_: "a"')
+    self.assertNotEqual('string_: "b"', 'string_: "a"')
+    self.assertNotEqual('string_: "ab"', 'string_: "aa"')
+
+    self.assertNotEqual('int64_: 0', '')
+    self.assertEqual('int64_: 0', 'int64_: 0')
+    self.assertNotEqual('int64_: -1', '')
+    self.assertNotEqual('int64_: 1', 'int64_: 0')
+    self.assertNotEqual('int64_: 0', 'int64_: -1')
+
+    self.assertNotEqual('float_: 0.0', '')
+    self.assertEqual('float_: 0.0', 'float_: 0.0')
+    self.assertNotEqual('float_: -0.1', '')
+    self.assertNotEqual('float_: 3.14', 'float_: 0')
+    self.assertNotEqual('float_: 0', 'float_: -0.1')
+    self.assertEqual('float_: -0.1', 'float_: -0.1')
+
+    self.assertNotEqual('bool_: true', '')
+    self.assertNotEqual('bool_: false', '')
+    self.assertNotEqual('bool_: true', 'bool_: false')
+    self.assertEqual('bool_: false', 'bool_: false')
+    self.assertEqual('bool_: true', 'bool_: true')
+
+    self.assertNotEqual('enum_: A', '')
+    self.assertNotEqual('enum_: B', 'enum_: A')
+    self.assertNotEqual('enum_: C', 'enum_: B')
+    self.assertEqual('enum_: C', 'enum_: C')
+
+  def testRepeatedPrimitives(self):
+    self.assertNotEqual('int64s: 0', '')
+    self.assertEqual('int64s: 0', 'int64s: 0')
+    self.assertNotEqual('int64s: 1', 'int64s: 0')
+    self.assertNotEqual('int64s: 0 int64s: 0', '')
+    self.assertNotEqual('int64s: 0 int64s: 0', 'int64s: 0')
+    self.assertNotEqual('int64s: 1 int64s: 0', 'int64s: 0')
+    self.assertNotEqual('int64s: 0 int64s: 1', 'int64s: 0')
+    self.assertNotEqual('int64s: 1', 'int64s: 0 int64s: 2')
+    self.assertNotEqual('int64s: 2 int64s: 0', 'int64s: 1')
+    self.assertEqual('int64s: 0 int64s: 0', 'int64s: 0 int64s: 0')
+    self.assertEqual('int64s: 0 int64s: 1', 'int64s: 0 int64s: 1')
+    self.assertNotEqual('int64s: 1 int64s: 0', 'int64s: 0 int64s: 0')
+    self.assertNotEqual('int64s: 1 int64s: 0', 'int64s: 0 int64s: 1')
+    self.assertNotEqual('int64s: 1 int64s: 0', 'int64s: 0 int64s: 2')
+    self.assertNotEqual('int64s: 1 int64s: 1', 'int64s: 1 int64s: 0')
+    self.assertNotEqual('int64s: 1 int64s: 1', 'int64s: 1 int64s: 0 int64s: 2')
+
+  def testMessage(self):
+    self.assertNotEqual('small <>', '')
+    self.assertEqual('small <>', 'small <>')
+    self.assertNotEqual('small < strings: "a" >', '')
+    self.assertNotEqual('small < strings: "a" >', 'small <>')
+    self.assertEqual('small < strings: "a" >', 'small < strings: "a" >')
+    self.assertNotEqual('small < strings: "b" >', 'small < strings: "a" >')
+    self.assertNotEqual(
+        'small < strings: "a" strings: "b" >', 'small < strings: "a" >'
+    )
+
+    self.assertNotEqual('string_: "a"', 'small <>')
+    self.assertNotEqual('string_: "a"', 'small < strings: "b" >')
+    self.assertNotEqual('string_: "a"', 'small < strings: "b" strings: "c" >')
+    self.assertNotEqual('string_: "a" small <>', 'small <>')
+    self.assertNotEqual('string_: "a" small <>', 'small < strings: "b" >')
+    self.assertEqual('string_: "a" small <>', 'string_: "a" small <>')
+    self.assertNotEqual(
+        'string_: "a" small < strings: "a" >', 'string_: "a" small <>'
+    )
+    self.assertEqual('string_: "a" small < strings: "a" >',
+                     'string_: "a" small < strings: "a" >')
+    self.assertNotEqual(
+        'string_: "a" small < strings: "a" >',
+        'int64_: 1 small < strings: "a" >',
+    )
+    self.assertNotEqual('string_: "a" small < strings: "a" >', 'int64_: 1')
+    self.assertNotEqual('string_: "a"', 'int64_: 1 small < strings: "a" >')
+    self.assertNotEqual(
+        'string_: "a" int64_: 0 small < strings: "a" >',
+        'int64_: 1 small < strings: "a" >',
+    )
+    self.assertNotEqual(
+        'string_: "a" int64_: 1 small < strings: "a" >',
+        'string_: "a" int64_: 0 small < strings: "a" >',
+    )
+    self.assertEqual('string_: "a" int64_: 0 small < strings: "a" >',
+                     'string_: "a" int64_: 0 small < strings: "a" >')
+
+  def testNestedMessage(self):
+    self.assertNotEqual('medium <>', '')
+    self.assertEqual('medium <>', 'medium <>')
+    self.assertNotEqual('medium < smalls <> >', 'medium <>')
+    self.assertEqual('medium < smalls <> >', 'medium < smalls <> >')
+    self.assertNotEqual(
+        'medium < smalls <> smalls <> >', 'medium < smalls <> >'
+    )
+    self.assertEqual('medium < smalls <> smalls <> >',
+                     'medium < smalls <> smalls <> >')
+
+    self.assertNotEqual('medium < int32s: 0 >', 'medium < smalls <> >')
+
+    self.assertNotEqual(
+        'medium < smalls < strings: "a"> >', 'medium < smalls <> >'
+    )
+
+  def testTagOrder(self):
+    """Tests that different fields are ordered by tag number.
+
+    For reference, here are the relevant tag numbers from compare_test.proto:
+      optional string string_ = 1;
+      optional int64 int64_ = 2;
+      optional float float_ = 3;
+      optional Small small = 8;
+      optional Medium medium = 7;
+      optional Small small = 8;
+    """
+    self.assertNotEqual(
+        'string_: "a"                      ',
+        '             int64_: 1            ',
+    )
+    self.assertNotEqual(
+        'string_: "a" int64_: 2            ',
+        '             int64_: 1            ',
+    )
+    self.assertNotEqual(
+        'string_: "b" int64_: 1            ',
+        'string_: "a" int64_: 2            ',
+    )
+    self.assertEqual('string_: "a" int64_: 1            ',
+                     'string_: "a" int64_: 1            ')
+    self.assertNotEqual(
+        'string_: "a" int64_: 1 float_: 0.0',
+        'string_: "a" int64_: 1            ',
+    )
+    self.assertEqual('string_: "a" int64_: 1 float_: 0.0',
+                     'string_: "a" int64_: 1 float_: 0.0')
+    self.assertNotEqual(
+        'string_: "a" int64_: 1 float_: 0.1',
+        'string_: "a" int64_: 1 float_: 0.0',
+    )
+    self.assertNotEqual(
+        'string_: "a" int64_: 2 float_: 0.0',
+        'string_: "a" int64_: 1 float_: 0.1',
+    )
+    self.assertNotEqual(
+        'string_: "a"                      ',
+        '             int64_: 1 float_: 0.1',
+    )
+    self.assertNotEqual(
+        'string_: "a"           float_: 0.0',
+        '             int64_: 1            ',
+    )
+    self.assertNotEqual(
+        'string_: "b"           float_: 0.0',
+        'string_: "a" int64_: 1            ',
+    )
+
+    self.assertNotEqual('string_: "a"', 'small < strings: "a" >')
+    self.assertNotEqual(
+        'string_: "a" small < strings: "a" >', 'small < strings: "b" >'
+    )
+    self.assertNotEqual(
+        'string_: "a" small < strings: "b" >',
+        'string_: "a" small < strings: "a" >',
+    )
+    self.assertEqual('string_: "a" small < strings: "a" >',
+                     'string_: "a" small < strings: "a" >')
+
+    self.assertNotEqual(
+        'string_: "a" medium <>', 'string_: "a" small < strings: "a" >'
+    )
+    self.assertNotEqual(
+        'string_: "a" medium < smalls <> >',
+        'string_: "a" small < strings: "a" >',
+    )
+    self.assertNotEqual('medium <>', 'small < strings: "a" >')
+    self.assertNotEqual('medium <> small <>', 'small < strings: "a" >')
+    self.assertNotEqual('medium < smalls <> >', 'small < strings: "a" >')
+    self.assertNotEqual(
+        'medium < smalls < strings: "a" > >', 'small < strings: "b" >'
+    )
+
+  def testIsClose(self):
+    self.assertTrue(compare.isClose(1, 1, 1e-10))
+    self.assertTrue(compare.isClose(65061.0420, 65061.0322, 1e-5))
+    self.assertFalse(compare.isClose(65061.0420, 65061.0322, 1e-7))
+
+  def testIsCloseNan(self):
+    self.assertTrue(compare.isClose(float('nan'), float('nan'), 1e-10))
+    self.assertFalse(compare.isClose(float('nan'), 1, 1e-10))
+    self.assertFalse(compare.isClose(1, float('nan'), 1e-10))
+    self.assertFalse(compare.isClose(float('nan'), float('inf'), 1e-10))
+
+  def testIsCloseInf(self):
+    self.assertTrue(compare.isClose(float('inf'), float('inf'), 1e-10))
+    self.assertTrue(compare.isClose(float('-inf'), float('-inf'), 1e-10))
+    self.assertFalse(compare.isClose(float('-inf'), float('inf'), 1e-10))
+    self.assertFalse(compare.isClose(float('inf'), 1, 1e-10))
+    self.assertFalse(compare.isClose(1, float('inf'), 1e-10))
+
+  def testIsCloseSubnormal(self):
+    x = sys.float_info.min * sys.float_info.epsilon
+    self.assertTrue(compare.isClose(x, x, 1e-10))
+    self.assertFalse(compare.isClose(x, 1, 1e-10))
+
+
+class NormalizeNumbersTest(googletest.TestCase):
+  """Tests for NormalizeNumberFields()."""
+
+  def testNormalizesInts(self):
+    pb = compare_test_pb2.Large(int64_=4)
+    compare.NormalizeNumberFields(pb)
+    self.assertIsInstance(pb.int64_, int)
+
+    pb.int64_ = 4
+    compare.NormalizeNumberFields(pb)
+    self.assertIsInstance(pb.int64_, int)
+
+    pb.int64_ = 9999999999999999
+    compare.NormalizeNumberFields(pb)
+    self.assertIsInstance(pb.int64_, int)
+
+  def testNormalizesRepeatedInts(self):
+    pb = compare_test_pb2.Large(int64s=[1, 400, 999999999999999])
+    compare.NormalizeNumberFields(pb)
+    self.assertIsInstance(pb.int64s[0], int)
+    self.assertIsInstance(pb.int64s[1], int)
+    self.assertIsInstance(pb.int64s[2], int)
+
+  def testNormalizesFloats(self):
+    pb1 = compare_test_pb2.Large(float_=1.2314352351231)
+    pb2 = compare_test_pb2.Large(float_=1.231435)
+    self.assertNotEqual(pb1.float_, pb2.float_)
+    compare.NormalizeNumberFields(pb1)
+    compare.NormalizeNumberFields(pb2)
+    self.assertEqual(pb1.float_, pb2.float_)
+
+  def testNormalizesRepeatedFloats(self):
+    pb = compare_test_pb2.Large(
+        medium=compare_test_pb2.Medium(floats=[0.111111111, 0.111111])
+    )
+    compare.NormalizeNumberFields(pb)
+    for value in pb.medium.floats:
+      self.assertAlmostEqual(0.111111, value)
+
+  def testNormalizesDoubles(self):
+    pb1 = compare_test_pb2.Large(double_=1.2314352351231)
+    pb2 = compare_test_pb2.Large(double_=1.2314352)
+    self.assertNotEqual(pb1.double_, pb2.double_)
+    compare.NormalizeNumberFields(pb1)
+    compare.NormalizeNumberFields(pb2)
+    self.assertEqual(pb1.double_, pb2.double_)
+
+  def testNormalizesMaps(self):
+    pb = compare_test_pb2.WithMap()
+    pb.value_message[4].strings.extend(['a', 'b', 'c'])
+    pb.value_string['d'] = 'e'
+    compare.NormalizeNumberFields(pb)
+
+
+class AssertTest(googletest.TestCase):
+  """Tests assertProtoEqual()."""
+
+  def assertProtoEqual(self, a, b, **kwargs):
+    if isinstance(a, str) and isinstance(b, str):
+      a, b = LargePbs(a, b)
+    compare.assertProtoEqual(self, a, b, **kwargs)
+
+  def assertAll(self, a, **kwargs):
+    """Checks that all possible asserts pass."""
+    self.assertProtoEqual(a, a, **kwargs)
+
+  def assertSameNotEqual(self, a, b):
+    """Checks that assertProtoEqual() fails."""
+    self.assertRaises(AssertionError, self.assertProtoEqual, a, b)
+
+  def assertNone(self, a, b, message, **kwargs):
+    """Checks that all possible asserts fail with the given message."""
+    message = re.escape(textwrap.dedent(message))
+    self.assertRaisesRegex(AssertionError, message, self.assertProtoEqual, a, b,
+                           **kwargs)
+
+  def testCheckInitialized(self):
+    # neither is initialized
+    a = compare_test_pb2.Labeled(optional=1)
+    self.assertNone(a, a, 'Initialization errors: ', check_initialized=True)
+    self.assertAll(a, check_initialized=False)
+
+    # a is initialized, b isn't
+    b = copy.deepcopy(a)
+    a.required = 2
+    self.assertNone(a, b, 'Initialization errors: ', check_initialized=True)
+    self.assertNone(
+        a,
+        b,
+        """
+                    - required: 2
+                      optional: 1
+                    """,
+        check_initialized=False)
+
+    # both are initialized
+    a = compare_test_pb2.Labeled(required=2)
+    self.assertAll(a, check_initialized=True)
+    self.assertAll(a, check_initialized=False)
+
+    b = copy.deepcopy(a)
+    b.required = 3
+    message = """
+              - required: 2
+              ?           ^
+              + required: 3
+              ?           ^
+              """
+    self.assertNone(a, b, message, check_initialized=True)
+    self.assertNone(a, b, message, check_initialized=False)
+
+  def testAssertEqualWithStringArg(self):
+    pb = compare_test_pb2.Large(string_='abc', float_=1.234)
+    compare.assertProtoEqual(self, """
+          string_: 'abc'
+          float_: 1.234
+        """, pb)
+
+  def testNormalizesNumbers(self):
+    pb1 = compare_test_pb2.Large(int64_=4)
+    pb2 = compare_test_pb2.Large(int64_=4)
+    compare.assertProtoEqual(self, pb1, pb2)
+
+  def testNormalizesFloat(self):
+    pb1 = compare_test_pb2.Large(double_=4.0)
+    pb2 = compare_test_pb2.Large(double_=4)
+    compare.assertProtoEqual(self, pb1, pb2, normalize_numbers=True)
+
+  def testLargeProtoData(self):
+    # Proto size should be larger than 2**16.
+    number_of_entries = 2**13
+    string_value = 'dummystr'  # Has length of 2**3.
+    pb1_txt = 'strings: "dummystr"\n' * number_of_entries
+    pb2 = compare_test_pb2.Small(strings=[string_value] * number_of_entries)
+    compare.assertProtoEqual(self, pb1_txt, pb2)
+
+    with self.assertRaises(AssertionError):
+      compare.assertProtoEqual(self, pb1_txt + 'strings: "Should fail."', pb2)
+
+  def testPrimitives(self):
+    self.assertAll('string_: "x"')
+    self.assertNone('string_: "x"', 'string_: "y"', """
+                    - string_: "x"
+                    ?           ^
+                    + string_: "y"
+                    ?           ^
+                    """)
+
+  def testRepeatedPrimitives(self):
+    self.assertAll('int64s: 0 int64s: 1')
+
+    self.assertSameNotEqual('int64s: 0 int64s: 1', 'int64s: 1 int64s: 0')
+    self.assertSameNotEqual('int64s: 0 int64s: 1 int64s: 2',
+                            'int64s: 2 int64s: 1 int64s: 0')
+
+    self.assertSameNotEqual('int64s: 0', 'int64s: 0 int64s: 0')
+    self.assertSameNotEqual('int64s: 0 int64s: 1',
+                            'int64s: 1 int64s: 0 int64s: 1')
+
+    self.assertNone('int64s: 0', 'int64s: 0 int64s: 2', """
+                      int64s: 0
+                    + int64s: 2
+                    """)
+    self.assertNone('int64s: 0 int64s: 1', 'int64s: 0 int64s: 2', """
+                      int64s: 0
+                    - int64s: 1
+                    ?         ^
+                    + int64s: 2
+                    ?         ^
+                    """)
+
+  def testMessage(self):
+    self.assertAll('medium: {}')
+    self.assertAll('medium: { smalls: {} }')
+    self.assertAll('medium: { int32s: 1 smalls: {} }')
+    self.assertAll('medium: { smalls: { strings: "x" } }')
+    self.assertAll(
+        'medium: { smalls: { strings: "x" } } small: { strings: "y" }')
+
+    self.assertSameNotEqual('medium: { smalls: { strings: "x" strings: "y" } }',
+                            'medium: { smalls: { strings: "y" strings: "x" } }')
+    self.assertSameNotEqual(
+        'medium: { smalls: { strings: "x" } smalls: { strings: "y" } }',
+        'medium: { smalls: { strings: "y" } smalls: { strings: "x" } }')
+
+    self.assertSameNotEqual(
+        'medium: { smalls: { strings: "x" strings: "y" strings: "x" } }',
+        'medium: { smalls: { strings: "y" strings: "x" } }')
+    self.assertSameNotEqual(
+        'medium: { smalls: { strings: "x" } int32s: 0 }',
+        'medium: { int32s: 0 smalls: { strings: "x" } int32s: 0 }')
+
+    self.assertNone('medium: {}', 'medium: { smalls: { strings: "x" } }', """
+                      medium {
+                    +   smalls {
+                    +     strings: "x"
+                    +   }
+                      }
+                    """)
+    self.assertNone('medium: { smalls: { strings: "x" } }',
+                    'medium: { smalls: {} }', """
+                      medium {
+                        smalls {
+                    -     strings: "x"
+                        }
+                      }
+                    """)
+    self.assertNone('medium: { int32s: 0 }', 'medium: { int32s: 1 }', """
+                      medium {
+                    -   int32s: 0
+                    ?           ^
+                    +   int32s: 1
+                    ?           ^
+                      }
+                    """)
+
+  def testMsgPassdown(self):
+    self.assertRaisesRegex(
+        AssertionError,
+        'test message passed down',
+        self.assertProtoEqual,
+        'medium: {}',
+        'medium: { smalls: { strings: "x" } }',
+        msg='test message passed down')
+
+  def testRepeatedMessage(self):
+    self.assertAll('medium: { smalls: {} smalls: {} }')
+    self.assertAll('medium: { smalls: { strings: "x" } } medium: {}')
+    self.assertAll('medium: { smalls: { strings: "x" } } medium: { int32s: 0 }')
+    self.assertAll('medium: { smalls: {} smalls: { strings: "x" } } small: {}')
+
+    self.assertSameNotEqual('medium: { smalls: { strings: "x" } smalls: {} }',
+                            'medium: { smalls: {} smalls: { strings: "x" } }')
+
+    self.assertSameNotEqual('medium: { smalls: {} }',
+                            'medium: { smalls: {} smalls: {} }')
+    self.assertSameNotEqual('medium: { smalls: {} smalls: {} } medium: {}',
+                            'medium: {} medium: {} medium: { smalls: {} }')
+    self.assertSameNotEqual(
+        'medium: { smalls: { strings: "x" } smalls: {} }',
+        'medium: { smalls: {} smalls: { strings: "x" } smalls: {} }')
+
+    self.assertNone('medium: {}', 'medium: {} medium { smalls: {} }', """
+                      medium {
+                    +   smalls {
+                    +   }
+                      }
+                    """)
+    self.assertNone('medium: { smalls: {} smalls: { strings: "x" } }',
+                    'medium: { smalls: {} smalls: { strings: "y" } }', """
+                      medium {
+                        smalls {
+                        }
+                        smalls {
+                    -     strings: "x"
+                    ?               ^
+                    +     strings: "y"
+                    ?               ^
+                        }
+                      }
+                    """)
+
+
+class MixinTests(compare.ProtoAssertions, googletest.TestCase):
+
+  def testAssertEqualWithStringArg(self):
+    pb = compare_test_pb2.Large(string_='abc', float_=1.234)
+    self.assertProtoEqual("""
+          string_: 'abc'
+          float_: 1.234
+        """, pb)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/util/py_checkpoint_reader_wrapper.cc

@@ -0,0 +1,151 @@
+/* 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.
+==============================================================================*/
+
+// Disallow Numpy 1.7 deprecated symbols.
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/types.h"
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+
+#include "numpy/arrayobject.h"
+#include "pybind11/chrono.h"  // from @pybind11
+#include "pybind11/complex.h"  // from @pybind11
+#include "pybind11/functional.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/c/checkpoint_reader.h"
+#include "tensorflow/c/safe_ptr.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/python/lib/core/ndarray_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;
+
+// TODO(amitpatankar): Move the custom type casters to separate common header
+// only libraries.
+
+namespace pybind11 {
+namespace detail {
+
+/* This is a custom type caster for the TensorShape object. For more
+ * documentation please refer to this link:
+ * https://pybind11.readthedocs.io/en/stable/advanced/cast/custom.html#custom-type-casters
+ * The PyCheckpointReader methods sometimes return the `TensorShape` object
+ * and the `DataType` object as outputs. This custom type caster helps Python
+ * handle it's conversion from C++ to Python. Since we do not accept these
+ * classes as arguments from Python, it is not necessary to define the `load`
+ * function to cast the object from Python to a C++ object.
+ */
+
+template <>
+struct type_caster<tensorflow::TensorShape> {
+ public:
+  PYBIND11_TYPE_CASTER(tensorflow::TensorShape, _("tensorflow::TensorShape"));
+
+  static handle cast(const tensorflow::TensorShape& src,
+                     return_value_policy unused_policy, handle unused_handle) {
+    // TODO(amitpatankar): Simplify handling TensorShape as output later.
+    size_t dims = src.dims();
+    tensorflow::Safe_PyObjectPtr value(PyList_New(dims));
+    for (size_t i = 0; i < dims; ++i) {
+#if PY_MAJOR_VERSION >= 3
+      tensorflow::Safe_PyObjectPtr dim_value(
+          tensorflow::make_safe(PyLong_FromLong(src.dim_size(i))));
+#else
+      tensorflow::Safe_PyObjectPtr dim_value(
+          tensorflow::make_safe(PyInt_FromLong(src.dim_size(i))));
+#endif
+      PyList_SET_ITEM(value.get(), i, dim_value.release());
+    }
+
+    return value.release();
+  }
+};
+
+template <>
+struct type_caster<tensorflow::DataType> {
+ public:
+  PYBIND11_TYPE_CASTER(tensorflow::DataType, _("tensorflow::DataType"));
+
+  static handle cast(const tensorflow::DataType& src,
+                     return_value_policy unused_policy, handle unused_handle) {
+#if PY_MAJOR_VERSION >= 3
+    tensorflow::Safe_PyObjectPtr value(
+        tensorflow::make_safe(PyLong_FromLong(src)));
+#else
+    tensorflow::Safe_PyObjectPtr value(
+        tensorflow::make_safe(PyInt_FromLong(src)));
+#endif
+    return value.release();
+  }
+};
+
+}  // namespace detail
+}  // namespace pybind11
+
+namespace tensorflow {
+
+static py::object CheckpointReader_GetTensor(
+    tensorflow::checkpoint::CheckpointReader* reader, const std::string& name) {
+  Safe_TF_StatusPtr status = make_safe(TF_NewStatus());
+  PyObject* py_obj = Py_None;
+  std::unique_ptr<tensorflow::Tensor> tensor;
+  reader->GetTensor(name, &tensor, status.get());
+
+  // Error handling if unable to get Tensor.
+  tensorflow::MaybeRaiseFromTFStatus(status.get());
+
+  tensorflow::MaybeRaiseFromStatus(
+      tensorflow::TensorToNdarray(*tensor, &py_obj));
+
+  return tensorflow::PyoOrThrow(
+      PyArray_Return(reinterpret_cast<PyArrayObject*>(py_obj)));
+}
+
+}  // namespace tensorflow
+
+PYBIND11_MODULE(_pywrap_checkpoint_reader, m) {
+  // Initialization code to use numpy types in the type casters.
+  import_array1();
+  py::class_<tensorflow::checkpoint::CheckpointReader> checkpoint_reader_class(
+      m, "CheckpointReader");
+  checkpoint_reader_class
+      .def(py::init([](const std::string& filename) {
+        tensorflow::Safe_TF_StatusPtr status =
+            tensorflow::make_safe(TF_NewStatus());
+        // pybind11 support smart pointers and will own freeing the memory when
+        // complete.
+        // https://pybind11.readthedocs.io/en/master/advanced/smart_ptrs.html#std-unique-ptr
+        auto checkpoint =
+            std::make_unique<tensorflow::checkpoint::CheckpointReader>(
+                filename, status.get());
+        tensorflow::MaybeRaiseFromTFStatus(status.get());
+        return checkpoint;
+      }))
+      .def("debug_string",
+           [](tensorflow::checkpoint::CheckpointReader& self) {
+             return py::bytes(self.DebugString());
+           })
+      .def("get_variable_to_shape_map",
+           &tensorflow::checkpoint::CheckpointReader::GetVariableToShapeMap)
+      .def("_GetVariableToDataTypeMap",
+           &tensorflow::checkpoint::CheckpointReader::GetVariableToDataTypeMap)
+      .def("_HasTensor", &tensorflow::checkpoint::CheckpointReader::HasTensor)
+      .def_static("CheckpointReader_GetTensor",
+                  &tensorflow::CheckpointReader_GetTensor);
+};

tensorflow/python/util/pywrap_xla_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 get_cpu_kernel_names() -> list[str]: ...
+def get_gpu_kernel_names() -> list[str]: ...

tensorflow/python/util/pywrap_xla_ops_test.py

@@ -0,0 +1,39 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from tensorflow.python.platform import googletest  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.util import pywrap_xla_ops
+
+
+class XlaOpsetUtilsTest(googletest.TestCase):
+
+  def testGetGpuCompilableKernelNames(self):
+    """Tests retrieving compilable op names for GPU."""
+    op_names = pywrap_xla_ops.get_gpu_kernel_names()
+    self.assertGreater(op_names.__len__(), 0)
+    self.assertEqual(op_names.count('Max'), 1)
+    self.assertEqual(op_names.count('Min'), 1)
+    self.assertEqual(op_names.count('MatMul'), 1)
+
+  def testGetCpuCompilableKernelNames(self):
+    """Tests retrieving compilable op names for CPU."""
+    op_names = pywrap_xla_ops.get_cpu_kernel_names()
+    self.assertGreater(op_names.__len__(), 0)
+    self.assertEqual(op_names.count('Max'), 1)
+    self.assertEqual(op_names.count('Min'), 1)
+    self.assertEqual(op_names.count('MatMul'), 1)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/util/serialization.py

@@ -0,0 +1,78 @@
+# 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.
+# ==============================================================================
+"""Utilities for serializing Python objects."""
+
+import numpy as np
+import wrapt
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.util.compat import collections_abc
+
+
+def get_json_type(obj):
+  """Serializes any object to a JSON-serializable structure.
+
+  Args:
+      obj: the object to serialize
+
+  Returns:
+      JSON-serializable structure representing `obj`.
+
+  Raises:
+      TypeError: if `obj` cannot be serialized.
+  """
+  # if obj is a serializable Keras class instance
+  # e.g. optimizer, layer
+  if hasattr(obj, 'get_config'):
+    return {'class_name': obj.__class__.__name__, 'config': obj.get_config()}
+
+  # if obj is any numpy type
+  if type(obj).__module__ == np.__name__:
+    if isinstance(obj, np.ndarray):
+      return obj.tolist()
+    else:
+      return obj.item()
+
+  # misc functions (e.g. loss function)
+  if callable(obj):
+    return obj.__name__
+
+  # if obj is a python 'type'
+  if type(obj).__name__ == type.__name__:
+    return obj.__name__
+
+  if isinstance(obj, tensor_shape.Dimension):
+    return obj.value
+
+  if isinstance(obj, tensor_shape.TensorShape):
+    return obj.as_list()
+
+  if isinstance(obj, dtypes.DType):
+    return obj.name
+
+  if isinstance(obj, collections_abc.Mapping):
+    return dict(obj)
+
+  if obj is Ellipsis:
+    return {'class_name': '__ellipsis__'}
+
+  if isinstance(obj, wrapt.ObjectProxy):
+    return obj.__wrapped__
+
+  raise TypeError(f'Object {obj} is not JSON-serializable. You may implement '
+                  'a `get_config()` method on the class '
+                  '(returning a JSON-serializable dictionary) to make it '
+                  'serializable.')

tensorflow/python/util/serialization_test.py

@@ -0,0 +1,35 @@
+# 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 serialization functions."""
+
+import json
+
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.platform import test
+from tensorflow.python.util import serialization
+
+
+class SerializationTests(test.TestCase):
+
+  def test_serialize_shape(self):
+    round_trip = json.loads(json.dumps(
+        tensor_shape.TensorShape([None, 2, 3]),
+        default=serialization.get_json_type))
+    self.assertIs(round_trip[0], None)
+    self.assertEqual(round_trip[1], 2)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/stack_trace.cc

@@ -0,0 +1,132 @@
+/* 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/python/util/stack_trace.h"
+
+#include <Python.h>
+
+#include <limits>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/base/attributes.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "tensorflow/core/platform/stack_frame.h"
+#include "tensorflow/core/util/managed_stack_trace.h"
+
+namespace {
+
+// Returns C string from a Python string object. Handles Python2/3 strings.
+// TODO(kkb): This is a generic Python utility function. factor out as a
+// utility.
+const char* GetPythonString(PyObject* o) {
+  if (PyBytes_Check(o)) {
+    return PyBytes_AsString(o);
+  } else {
+    return PyUnicode_AsUTF8(o);
+  }
+}
+
+}  // namespace
+
+namespace tensorflow {
+
+ABSL_MUST_USE_RESULT
+ABSL_ATTRIBUTE_HOT
+std::shared_ptr<StackTrace> StackTrace::Capture(int limit) {
+  DCHECK(PyGILState_Check());
+  if (limit == -1) limit = std::numeric_limits<int>::max();
+
+  StackTrace result;
+#if PY_VERSION_HEX >= 0x030B0000
+  PyFrameObject* oldframe;
+  PyFrameObject* frame = PyThreadState_GetFrame(PyThreadState_GET());
+  for (int i = 0; i < limit && frame != nullptr; oldframe = frame,
+           frame = PyFrame_GetBack(frame), Py_DECREF(oldframe), ++i) {
+    PyCodeObject* code_obj = PyFrame_GetCode(frame);
+    DCHECK(code_obj != nullptr);
+
+    int line_number = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(frame));
+    result.code_objs_.push_back(std::make_pair(code_obj, line_number));
+  }
+  Py_XDECREF(frame);
+#else
+  const PyFrameObject* frame = PyThreadState_GET()->frame;
+  for (int i = 0; i < limit && frame != nullptr; frame = frame->f_back, ++i) {
+    PyCodeObject* code_obj = frame->f_code;
+    Py_XINCREF(code_obj);
+    DCHECK(code_obj != nullptr);
+
+    int line_number = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(frame));
+    result.code_objs_.push_back(std::make_pair(code_obj, line_number));
+  }
+#endif
+
+  static absl::flat_hash_map<uint64_t, std::shared_ptr<StackTrace>>* cache =
+      new absl::flat_hash_map<uint64_t, std::shared_ptr<StackTrace>>();
+#ifdef Py_GIL_DISABLED
+  static absl::Mutex mu(absl::kConstInit);
+
+  absl::MutexLock lock(&mu);
+#endif  // Py_GIL_DISABLED
+  uint64_t hash_code = result.hash();
+  if (!cache->contains(hash_code)) {
+    cache->insert(std::make_pair(
+        hash_code, std::make_shared<StackTrace>(std::move(result))));
+  }
+
+  return cache->at(hash_code);
+}
+
+std::vector<StackFrame> StackTrace::ToStackFrames(
+    const SourceMap& source_map, const StackTraceFilter& filtered,
+    bool reverse_traversal, int limit) const {
+  auto gil_state = PyGILState_Ensure();
+  std::vector<StackFrame> result;
+  result.reserve(code_objs_.size());
+
+  if (limit == -1) limit = std::numeric_limits<int>::max();
+
+  for (int i = 0; i < code_objs_.size(); i++) {
+    int idx = reverse_traversal ? i : code_objs_.size() - 1 - i;
+
+    const std::pair<PyCodeObject*, int>& code_obj = code_objs_[idx];
+    const char* file_name = GetPythonString(code_obj.first->co_filename);
+    const int line_number = code_obj.second;
+
+    if (filtered && filtered(file_name)) {
+      continue;
+    }
+
+    const auto it = source_map.find(SourceLoc{file_name, line_number});
+    if (it != source_map.end()) {
+      result.push_back(it->second);
+    } else {
+      result.emplace_back(StackFrame{file_name, line_number,
+                                     GetPythonString(code_obj.first->co_name)});
+    }
+
+    if (result.size() == limit) {
+      break;
+    }
+  }
+
+  PyGILState_Release(gil_state);
+  return result;
+}
+
+}  // namespace tensorflow

tensorflow/python/util/stack_trace.h

@@ -0,0 +1,94 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
+#define TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
+
+#include <Python.h>
+#include <frameobject.h>
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/base/attributes.h"
+#include "absl/container/inlined_vector.h"
+#include "tensorflow/core/platform/stack_frame.h"
+#include "tensorflow/core/util/managed_stack_trace.h"
+#include "tsl/platform/fingerprint.h"
+
+namespace tensorflow {
+
+// A class for capturing Python stack trace.
+class StackTrace : public CapturedStackTrace {
+ public:
+  static constexpr int kStackTraceInitialSize = 30;
+
+  StackTrace() = default;
+  StackTrace(StackTrace&& other) = default;
+
+  // Returns `StackTrace` object that captures the current Python stack trace.
+  // `limit` determines how many stack frames at most are returned: set to -1
+  // for "no limit".
+  // Python GIL must be acquired beforehand.
+  ABSL_MUST_USE_RESULT
+  ABSL_ATTRIBUTE_HOT
+  static std::shared_ptr<StackTrace> Capture(int limit);
+
+  uint64_t hash() const {
+    uint64_t hash = 0;
+    for (const auto& p : code_objs_) {
+      hash = tsl::FingerprintCat64(hash,
+                                   reinterpret_cast<const uint64_t>(p.first));
+      hash = tsl::FingerprintCat64(hash, p.second);
+    }
+    return hash;
+  }
+
+  ~StackTrace() override { Clear(); }
+
+  // Returns a structured representation of the captured stack trace.
+  // `source_map` provides a custom mapping for translating stack frames,
+  // `filter` returns `true` for the stack frames which should be omitted.
+  //
+  // `reverse_traversal` changes the traversal order of the stack trace, and
+  // `limit` bounds the number of returned frames (after filtering).
+  std::vector<StackFrame> ToStackFrames(const SourceMap& source_map,
+                                        const StackTraceFilter& filtered,
+                                        bool reverse_traversal,
+                                        int limit) const override;
+
+ private:
+  void Clear() {
+    auto gil_state = PyGILState_Ensure();
+    for (const auto& p : code_objs_) Py_DECREF(p.first);
+    code_objs_.clear();
+    PyGILState_Release(gil_state);
+  }
+
+  absl::InlinedVector<std::pair<PyCodeObject*, int>, 16> code_objs_;
+
+  StackTrace(const StackTrace&) = delete;
+  StackTrace& operator=(const StackTrace&) = delete;
+};
+
+inline std::shared_ptr<StackTrace> GetStackTrace(int limit = -1) {
+  return StackTrace::Capture(limit);
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_

tensorflow/python/util/stat_summarizer_wrapper.cc

@@ -0,0 +1,48 @@
+/* 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.
+==============================================================================*/
+
+#include <string>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "xla/tsl/util/stat_summarizer_options.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/step_stats.pb.h"
+#include "tensorflow/core/util/stat_summarizer.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_stat_summarizer, m, pybind11::mod_gil_not_used()) {
+  py::class_<tensorflow::StatSummarizer> stat_summ_class(m, "StatSummarizer",
+                                                         py::dynamic_attr());
+  stat_summ_class
+      .def(py::init([](std::string graph_def_serialized) {
+        tensorflow::GraphDef proto;
+        proto.ParseFromString(graph_def_serialized);
+        return new tensorflow::StatSummarizer(proto);
+      }))
+      .def(py::init([]() {
+        return new tensorflow::StatSummarizer(tsl::StatSummarizerOptions());
+      }))
+      .def("ProcessStepStats", &tensorflow::StatSummarizer::ProcessStepStats)
+      .def("GetOutputString", &tensorflow::StatSummarizer::GetOutputString)
+      .def("PrintStepStats", &tensorflow::StatSummarizer::PrintStepStats)
+      .def("ProcessStepStatsStr", [](tensorflow::StatSummarizer& self,
+                                     const std::string& step_stats_str) {
+        tensorflow::StepStats step_stats;
+        step_stats.ParseFromString(step_stats_str);
+        self.ProcessStepStats(step_stats);
+      });
+};

tensorflow/python/util/tensor_float_32.cc

@@ -0,0 +1,22 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/platform/tensor_float_32_utils.h"
+
+PYBIND11_MODULE(_pywrap_tensor_float_32_execution, m) {
+  m.def("enable", &tensorflow::enable_tensor_float_32_execution);
+  m.def("is_enabled", &tensorflow::tensor_float_32_execution_enabled);
+}

tensorflow/python/util/tf2xla_opset_wrapper.cc

@@ -0,0 +1,49 @@
+/* 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.
+==============================================================================*/
+
+#include <pybind11/stl.h>
+
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil
+#include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil
+#include "tensorflow/compiler/tf2xla/tf2xla_opset.h"
+
+using tensorflow::GetRegisteredXlaOpsForDevice;
+
+PYBIND11_MODULE(pywrap_xla_ops, m) {
+  pybind11::google::ImportStatusModule();
+  m.def(
+      "get_gpu_kernel_names",
+      []() -> absl::StatusOr<std::vector<std::string>> {
+        return GetRegisteredXlaOpsForDevice("XLA_GPU_JIT");
+      },
+      R"pbdoc(
+     Returns list of names of gpu ops that can be compiled.
+    )pbdoc");
+  m.def(
+      "get_cpu_kernel_names",
+      []() -> absl::StatusOr<std::vector<std::string>> {
+        return GetRegisteredXlaOpsForDevice("XLA_CPU_JIT");
+      },
+      R"pbdoc(
+     Returns list of names of cpu ops that can be compiled.
+    )pbdoc");
+};

tensorflow/python/util/tf_contextlib.py

@@ -0,0 +1,39 @@
+# 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.
+# ==============================================================================
+"""TFDecorator-aware replacements for the contextlib module."""
+from collections.abc import Callable, Iterator
+import contextlib as _contextlib
+
+from typing import ContextManager, TypeVar
+
+from tensorflow.python.util import tf_decorator
+
+_T = TypeVar('_T')
+
+
+def contextmanager(
+    target: Callable[..., Iterator[_T]],
+) -> Callable[..., ContextManager[_T]]:
+  """A tf_decorator-aware wrapper for `contextlib.contextmanager`.
+
+  Usage is identical to `contextlib.contextmanager`.
+
+  Args:
+    target: A callable to be wrapped in a contextmanager.
+  Returns:
+    A callable that can be used inside of a `with` statement.
+  """
+  context_manager = _contextlib.contextmanager(target)
+  return tf_decorator.make_decorator(target, context_manager, 'contextmanager')

tensorflow/python/util/tf_contextlib_test.py

@@ -0,0 +1,88 @@
+# 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.
+# ==============================================================================
+"""Unit tests for tf_contextlib."""
+
+# pylint: disable=unused-import
+from tensorflow.python.platform import test
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+
+
+@tf_contextlib.contextmanager
+def test_yield_append_before_and_after_yield(x, before, after):
+  x.append(before)
+  yield
+  x.append(after)
+
+
+@tf_contextlib.contextmanager
+def test_yield_return_x_plus_1(x):
+  yield x + 1
+
+
+@tf_contextlib.contextmanager
+def test_params_and_defaults(a, b=2, c=True, d='hello'):
+  return [a, b, c, d]
+
+
+class TfContextlibTest(test.TestCase):
+
+  def testRunsCodeBeforeYield(self):
+    x = []
+    with test_yield_append_before_and_after_yield(x, 'before', ''):
+      self.assertEqual('before', x[-1])
+
+  def testRunsCodeAfterYield(self):
+    x = []
+    with test_yield_append_before_and_after_yield(x, '', 'after'):
+      pass
+    self.assertEqual('after', x[-1])
+
+  def testNestedWith(self):
+    x = []
+    with test_yield_append_before_and_after_yield(x, 'before', 'after'):
+      with test_yield_append_before_and_after_yield(x, 'inner', 'outer'):
+        with test_yield_return_x_plus_1(1) as var:
+          x.append(var)
+    self.assertEqual(['before', 'inner', 2, 'outer', 'after'], x)
+
+  def testMultipleCallsOfSeparateInstances(self):
+    x = []
+    with test_yield_append_before_and_after_yield(x, 1, 2):
+      pass
+    with test_yield_append_before_and_after_yield(x, 3, 4):
+      pass
+    self.assertEqual([1, 2, 3, 4], x)
+
+  def testReturnsResultFromYield(self):
+    with test_yield_return_x_plus_1(3) as result:
+      self.assertEqual(4, result)
+
+  def testUnwrapContextManager(self):
+    decorators, target = tf_decorator.unwrap(test_params_and_defaults)
+    self.assertEqual(1, len(decorators))
+    self.assertTrue(isinstance(decorators[0], tf_decorator.TFDecorator))
+    self.assertEqual('contextmanager', decorators[0].decorator_name)
+    self.assertFalse(isinstance(target, tf_decorator.TFDecorator))
+
+  def testGetArgSpecReturnsWrappedArgSpec(self):
+    argspec = tf_inspect.getargspec(test_params_and_defaults)
+    self.assertEqual(['a', 'b', 'c', 'd'], argspec.args)
+    self.assertEqual((2, True, 'hello'), argspec.defaults)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/tf_decorator.py

@@ -0,0 +1,361 @@
+# 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.
+# ==============================================================================
+"""Base TFDecorator class and utility functions for working with decorators.
+
+There are two ways to create decorators that TensorFlow can introspect into.
+This is important for documentation generation purposes, so that function
+signatures aren't obscured by the (*args, **kwds) signature that decorators
+often provide.
+
+1. Call `tf_decorator.make_decorator` on your wrapper function. If your
+decorator is stateless, or can capture all of the variables it needs to work
+with through lexical closure, this is the simplest option. Create your wrapper
+function as usual, but instead of returning it, return
+`tf_decorator.make_decorator(target, your_wrapper)`. This will attach some
+decorator introspection metadata onto your wrapper and return it.
+
+Example:
+
+  def print_hello_before_calling(target):
+    def wrapper(*args, **kwargs):
+      print('hello')
+      return target(*args, **kwargs)
+    return tf_decorator.make_decorator(target, wrapper)
+
+2. Derive from TFDecorator. If your decorator needs to be stateful, you can
+implement it in terms of a TFDecorator. Store whatever state you need in your
+derived class, and implement the `__call__` method to do your work before
+calling into your target. You can retrieve the target via
+`super(MyDecoratorClass, self).decorated_target`, and call it with whatever
+parameters it needs.
+
+Example:
+
+  class CallCounter(tf_decorator.TFDecorator):
+    def __init__(self, target):
+      super(CallCounter, self).__init__('count_calls', target)
+      self.call_count = 0
+
+    def __call__(self, *args, **kwargs):
+      self.call_count += 1
+      return super(CallCounter, self).decorated_target(*args, **kwargs)
+
+  def count_calls(target):
+    return CallCounter(target)
+"""
+import inspect
+from typing import Dict, Any
+
+
+def _make_default_values(fullargspec: inspect.FullArgSpec) -> Dict[str, Any]:
+  """Returns default values from the function's fullargspec."""
+  if fullargspec.defaults is not None:
+    defaults = {
+        name: value for name, value in zip(
+            fullargspec.args[-len(fullargspec.defaults):], fullargspec.defaults)
+    }
+  else:
+    defaults = {}
+
+  if fullargspec.kwonlydefaults is not None:
+    defaults.update(fullargspec.kwonlydefaults)
+
+  return defaults
+
+
+def fullargspec_to_signature(
+    fullargspec: inspect.FullArgSpec) -> inspect.Signature:
+  """Repackages fullargspec information into an equivalent inspect.Signature."""
+  defaults = _make_default_values(fullargspec)
+  parameters = []
+
+  for arg in fullargspec.args:
+    parameters.append(
+        inspect.Parameter(
+            arg,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=defaults.get(arg, inspect.Parameter.empty),
+        )
+    )
+
+  if fullargspec.varargs is not None:
+    parameters.append(
+        inspect.Parameter(fullargspec.varargs, inspect.Parameter.VAR_POSITIONAL)
+    )
+
+  for kwarg in fullargspec.kwonlyargs:
+    parameters.append(
+        inspect.Parameter(
+            kwarg,
+            inspect.Parameter.KEYWORD_ONLY,
+            default=defaults.get(kwarg, inspect.Parameter.empty),
+        )
+    )
+
+  if fullargspec.varkw is not None:
+    parameters.append(
+        inspect.Parameter(fullargspec.varkw, inspect.Parameter.VAR_KEYWORD)
+    )
+
+  return inspect.Signature(parameters)
+
+
+def make_decorator(target,
+                   decorator_func,
+                   decorator_name=None,
+                   decorator_doc='',
+                   decorator_argspec=None):
+  """Make a decorator from a wrapper and a target.
+
+  Args:
+    target: The final callable to be wrapped.
+    decorator_func: The wrapper function.
+    decorator_name: The name of the decorator. If `None`, the name of the
+      function calling make_decorator.
+    decorator_doc: Documentation specific to this application of
+      `decorator_func` to `target`.
+    decorator_argspec: Override the signature using FullArgSpec.
+
+  Returns:
+    The `decorator_func` argument with new metadata attached.
+  """
+  if decorator_name is None:
+    decorator_name = inspect.currentframe().f_back.f_code.co_name
+  decorator = TFDecorator(decorator_name, target, decorator_doc,
+                          decorator_argspec)
+  setattr(decorator_func, '_tf_decorator', decorator)
+  # Objects that are callables (e.g., a functools.partial object) may not have
+  # the following attributes.
+  if hasattr(target, '__name__'):
+    decorator_func.__name__ = target.__name__
+  if hasattr(target, '__qualname__'):
+    decorator_func.__qualname__ = target.__qualname__
+  if hasattr(target, '__module__'):
+    decorator_func.__module__ = target.__module__
+  if hasattr(target, '__dict__'):
+    # Copy dict entries from target which are not overridden by decorator_func.
+    for name in target.__dict__:
+      if name not in decorator_func.__dict__:
+        decorator_func.__dict__[name] = target.__dict__[name]
+  if hasattr(target, '__doc__'):
+    decorator_func.__doc__ = decorator.__doc__
+  decorator_func.__wrapped__ = target
+  # Keeping a second handle to `target` allows callers to detect whether the
+  # decorator was modified using `rewrap`.
+  decorator_func.__original_wrapped__ = target
+  if decorator_argspec:
+    decorator_func.__signature__ = fullargspec_to_signature(
+        decorator_argspec)
+  elif callable(target):
+    try:
+      signature = inspect.signature(target)
+    except (TypeError, ValueError):
+      # Certain callables such as builtins can not be inspected for signature.
+      pass
+    else:
+      bound_instance = _get_bound_instance(target)
+      # Present the decorated func as a method as well
+      if bound_instance and 'self' in signature.parameters:
+        signature = inspect.Signature(list(signature.parameters.values())[1:])
+        decorator_func.__self__ = bound_instance
+
+      decorator_func.__signature__ = signature
+
+  return decorator_func
+
+
+def _get_bound_instance(target):
+  """Returns the instance any of the targets is attached to."""
+  decorators, target = unwrap(target)
+  for decorator in decorators:
+    if inspect.ismethod(decorator.decorated_target):
+      return decorator.decorated_target.__self__
+
+
+def _has_tf_decorator_attr(obj):
+  """Checks if object has _tf_decorator attribute.
+
+  This check would work for mocked object as well since it would
+  check if returned attribute has the right type.
+
+  Args:
+    obj: Python object.
+  """
+  return (hasattr(obj, '_tf_decorator') and
+          isinstance(getattr(obj, '_tf_decorator'), TFDecorator))
+
+
+def rewrap(decorator_func, previous_target, new_target):
+  """Injects a new target into a function built by make_decorator.
+
+  This function allows replacing a function wrapped by `decorator_func`,
+  assuming the decorator that wraps the function is written as described below.
+
+  The decorator function must use `<decorator name>.__wrapped__` instead of the
+  wrapped function that is normally used:
+
+  Example:
+
+      # Instead of this:
+      def simple_parametrized_wrapper(*args, **kwds):
+        return wrapped_fn(*args, **kwds)
+
+      tf_decorator.make_decorator(simple_parametrized_wrapper, wrapped_fn)
+
+      # Write this:
+      def simple_parametrized_wrapper(*args, **kwds):
+        return simple_parametrized_wrapper.__wrapped__(*args, **kwds)
+
+      tf_decorator.make_decorator(simple_parametrized_wrapper, wrapped_fn)
+
+  Note that this process modifies decorator_func.
+
+  Args:
+    decorator_func: Callable returned by `wrap`.
+    previous_target: Callable that needs to be replaced.
+    new_target: Callable to replace previous_target with.
+
+  Returns:
+    The updated decorator. If decorator_func is not a tf_decorator, new_target
+    is returned.
+  """
+  # Because the process mutates the decorator, we only need to alter the
+  # innermost function that wraps previous_target.
+  cur = decorator_func
+  innermost_decorator = None
+  target = None
+  while _has_tf_decorator_attr(cur):
+    innermost_decorator = cur
+    target = getattr(cur, '_tf_decorator')
+    if target.decorated_target is previous_target:
+      break
+    cur = target.decorated_target
+    assert cur is not None
+
+  # If decorator_func is not a decorator, new_target replaces it directly.
+  if innermost_decorator is None:
+    # Consistency check. The caller should always pass the result of
+    # tf_decorator.unwrap as previous_target. If decorator_func is not a
+    # decorator, that will have returned decorator_func itself.
+    assert decorator_func is previous_target
+    return new_target
+
+  target.decorated_target = new_target
+
+  if inspect.ismethod(innermost_decorator):
+    # Bound methods can't be assigned attributes. Thankfully, they seem to
+    # be just proxies for their unbound counterpart, and we can modify that.
+    if hasattr(innermost_decorator, '__func__'):
+      innermost_decorator.__func__.__wrapped__ = new_target
+    elif hasattr(innermost_decorator, 'im_func'):
+      innermost_decorator.im_func.__wrapped__ = new_target
+    else:
+      innermost_decorator.__wrapped__ = new_target
+  else:
+    innermost_decorator.__wrapped__ = new_target
+
+  return decorator_func
+
+
+def unwrap(maybe_tf_decorator):
+  """Unwraps an object into a list of TFDecorators and a final target.
+
+  Args:
+    maybe_tf_decorator: Any callable object.
+
+  Returns:
+    A tuple whose first element is an list of TFDecorator-derived objects that
+    were applied to the final callable target, and whose second element is the
+    final undecorated callable target. If the `maybe_tf_decorator` parameter is
+    not decorated by any TFDecorators, the first tuple element will be an empty
+    list. The `TFDecorator` list is ordered from outermost to innermost
+    decorators.
+  """
+  decorators = []
+  cur = maybe_tf_decorator
+  while True:
+    if isinstance(cur, TFDecorator):
+      decorators.append(cur)
+    elif _has_tf_decorator_attr(cur):
+      decorators.append(getattr(cur, '_tf_decorator'))
+    else:
+      break
+    if not hasattr(decorators[-1], 'decorated_target'):
+      break
+    cur = decorators[-1].decorated_target
+  return decorators, cur
+
+
+class TFDecorator(object):
+  """Base class for all TensorFlow decorators.
+
+  TFDecorator captures and exposes the wrapped target, and provides details
+  about the current decorator.
+  """
+
+  def __init__(self,
+               decorator_name,
+               target,
+               decorator_doc='',
+               decorator_argspec=None):
+    self._decorated_target = target
+    self._decorator_name = decorator_name
+    self._decorator_doc = decorator_doc
+    self._decorator_argspec = decorator_argspec
+    if hasattr(target, '__name__'):
+      self.__name__ = target.__name__
+    if hasattr(target, '__qualname__'):
+      self.__qualname__ = target.__qualname__
+    if self._decorator_doc:
+      self.__doc__ = self._decorator_doc
+    elif hasattr(target, '__doc__') and target.__doc__:
+      self.__doc__ = target.__doc__
+    else:
+      self.__doc__ = ''
+
+    if decorator_argspec:
+      self.__signature__ = fullargspec_to_signature(decorator_argspec)
+    elif callable(target):
+      try:
+        self.__signature__ = inspect.signature(target)
+      except (TypeError, ValueError):
+        # Certain callables such as builtins can not be inspected for signature.
+        pass
+
+  def __get__(self, instance, owner):
+    return self._decorated_target.__get__(instance, owner)
+
+  def __call__(self, *args, **kwargs):
+    return self._decorated_target(*args, **kwargs)
+
+  @property
+  def decorated_target(self):
+    return self._decorated_target
+
+  @decorated_target.setter
+  def decorated_target(self, decorated_target):
+    self._decorated_target = decorated_target
+
+  @property
+  def decorator_name(self):
+    return self._decorator_name
+
+  @property
+  def decorator_doc(self):
+    return self._decorator_doc
+
+  @property
+  def decorator_argspec(self):
+    return self._decorator_argspec

tensorflow/python/util/tf_decorator_export.py

@@ -0,0 +1,26 @@
+# 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.
+# ==============================================================================
+"""Exports functions from tf_decorator.py to avoid cycles."""
+
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_export
+
+
+make_decorator = tf_export.tf_export(
+    '__internal__.decorator.make_decorator', v1=[]
+)(tf_decorator.make_decorator)
+unwrap = tf_export.tf_export('__internal__.decorator.unwrap', v1=[])(
+    tf_decorator.unwrap
+)

tensorflow/python/util/tf_decorator_test.py

@@ -0,0 +1,371 @@
+# 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.
+# ==============================================================================
+"""Unit tests for tf_decorator."""
+
+# pylint: disable=unused-import
+import functools
+import inspect
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+
+
+def test_tfdecorator(decorator_name, decorator_doc=None):
+
+  def make_tf_decorator(target):
+    return tf_decorator.TFDecorator(decorator_name, target, decorator_doc)
+
+  return make_tf_decorator
+
+
+def test_decorator_increment_first_int_arg(target):
+  """This test decorator skips past `self` as args[0] in the bound case."""
+
+  def wrapper(*args, **kwargs):
+    new_args = []
+    found = False
+    for arg in args:
+      if not found and isinstance(arg, int):
+        new_args.append(arg + 1)
+        found = True
+      else:
+        new_args.append(arg)
+    return target(*new_args, **kwargs)
+
+  return tf_decorator.make_decorator(target, wrapper)
+
+
+def test_injectable_decorator_square(target):
+
+  def wrapper(x):
+    return wrapper.__wrapped__(x)**2
+
+  return tf_decorator.make_decorator(target, wrapper)
+
+
+def test_injectable_decorator_increment(target):
+
+  def wrapper(x):
+    return wrapper.__wrapped__(x) + 1
+
+  return tf_decorator.make_decorator(target, wrapper)
+
+
+def test_function(x):
+  """Test Function Docstring."""
+  return x + 1
+
+
+@test_tfdecorator('decorator 1')
+@test_decorator_increment_first_int_arg
+@test_tfdecorator('decorator 3', 'decorator 3 documentation')
+def test_decorated_function(x):
+  """Test Decorated Function Docstring."""
+  return x * 2
+
+
+@test_tfdecorator('decorator')
+class TestDecoratedClass(object):
+  """Test Decorated Class."""
+
+  def __init__(self, two_attr=2):
+    self.two_attr = two_attr
+
+  @property
+  def two_prop(self):
+    return 2
+
+  def two_func(self):
+    return 2
+
+  @test_decorator_increment_first_int_arg
+  def return_params(self, a, b, c):
+    """Return parameters."""
+    return [a, b, c]
+
+
+class TfDecoratorTest(test.TestCase):
+
+  def testInitCapturesTarget(self):
+    self.assertIs(test_function,
+                  tf_decorator.TFDecorator('', test_function).decorated_target)
+
+  def testInitCapturesDecoratorName(self):
+    self.assertEqual(
+        'decorator name',
+        tf_decorator.TFDecorator('decorator name',
+                                 test_function).decorator_name)
+
+  def testInitCapturesDecoratorDoc(self):
+    self.assertEqual(
+        'decorator doc',
+        tf_decorator.TFDecorator('', test_function,
+                                 'decorator doc').decorator_doc)
+
+  def testInitCapturesNonNoneArgspec(self):
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations=None,
+    )
+    self.assertIs(
+        argspec,
+        tf_decorator.TFDecorator('', test_function, '',
+                                 argspec).decorator_argspec)
+
+  def testInitSetsDecoratorNameToTargetName(self):
+    self.assertEqual('test_function',
+                     tf_decorator.TFDecorator('', test_function).__name__)
+
+  def testInitSetsDecoratorQualNameToTargetQualName(self):
+    if hasattr(tf_decorator.TFDecorator('', test_function), '__qualname__'):
+      self.assertEqual('test_function',
+                       tf_decorator.TFDecorator('', test_function).__qualname__)
+
+  def testInitSetsDecoratorDocToTargetDoc(self):
+    self.assertEqual('Test Function Docstring.',
+                     tf_decorator.TFDecorator('', test_function).__doc__)
+
+  def testCallingATFDecoratorCallsTheTarget(self):
+    self.assertEqual(124, tf_decorator.TFDecorator('', test_function)(123))
+
+  def testCallingADecoratedFunctionCallsTheTarget(self):
+    self.assertEqual((2 + 1) * 2, test_decorated_function(2))
+
+  def testInitializingDecoratedClassWithInitParamsDoesntRaise(self):
+    try:
+      TestDecoratedClass(2)
+    except TypeError:
+      self.assertFail()
+
+  def testReadingClassAttributeOnDecoratedClass(self):
+    self.assertEqual(2, TestDecoratedClass().two_attr)
+
+  def testCallingClassMethodOnDecoratedClass(self):
+    self.assertEqual(2, TestDecoratedClass().two_func())
+
+  def testReadingClassPropertyOnDecoratedClass(self):
+    self.assertEqual(2, TestDecoratedClass().two_prop)
+
+  def testNameOnBoundProperty(self):
+    self.assertEqual('return_params',
+                     TestDecoratedClass().return_params.__name__)
+
+  def testQualNameOnBoundProperty(self):
+    if hasattr(TestDecoratedClass().return_params, '__qualname__'):
+      self.assertEqual('TestDecoratedClass.return_params',
+                       TestDecoratedClass().return_params.__qualname__)
+
+  def testDocstringOnBoundProperty(self):
+    self.assertEqual('Return parameters.',
+                     TestDecoratedClass().return_params.__doc__)
+
+  def testTarget__get__IsProxied(self):
+
+    class Descr(object):
+
+      def __get__(self, instance, owner):
+        return self
+
+    class Foo(object):
+      foo = tf_decorator.TFDecorator('Descr', Descr())
+
+    self.assertIsInstance(Foo.foo, Descr)
+
+
+def test_wrapper(*args, **kwargs):
+  return test_function(*args, **kwargs)
+
+
+class TfMakeDecoratorTest(test.TestCase):
+
+  def testAttachesATFDecoratorAttr(self):
+    decorated = tf_decorator.make_decorator(test_function, test_wrapper)
+    decorator = getattr(decorated, '_tf_decorator')
+    self.assertIsInstance(decorator, tf_decorator.TFDecorator)
+
+  def testAttachesWrappedAttr(self):
+    decorated = tf_decorator.make_decorator(test_function, test_wrapper)
+    wrapped_attr = getattr(decorated, '__wrapped__')
+    self.assertIs(test_function, wrapped_attr)
+
+  def testSetsTFDecoratorNameToDecoratorNameArg(self):
+    decorated = tf_decorator.make_decorator(test_function, test_wrapper,
+                                            'test decorator name')
+    decorator = getattr(decorated, '_tf_decorator')
+    self.assertEqual('test decorator name', decorator.decorator_name)
+
+  def testSetsTFDecoratorDocToDecoratorDocArg(self):
+    decorated = tf_decorator.make_decorator(
+        test_function, test_wrapper, decorator_doc='test decorator doc')
+    decorator = getattr(decorated, '_tf_decorator')
+    self.assertEqual('test decorator doc', decorator.decorator_doc)
+
+  def testUpdatesDictWithMissingEntries(self):
+    test_function.foobar = True
+    decorated = tf_decorator.make_decorator(test_function, test_wrapper)
+    self.assertTrue(decorated.foobar)
+    del test_function.foobar
+
+  def testUpdatesDict_doesNotOverridePresentEntries(self):
+    test_function.foobar = True
+    test_wrapper.foobar = False
+    decorated = tf_decorator.make_decorator(test_function, test_wrapper)
+    self.assertFalse(decorated.foobar)
+    del test_function.foobar
+    del test_wrapper.foobar
+
+  def testSetsTFDecoratorArgSpec(self):
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs='args',
+        kwonlyargs={},
+        defaults=(1, 'hello'),
+        kwonlydefaults=None,
+        varkw='kwargs',
+        annotations=None)
+    decorated = tf_decorator.make_decorator(test_function, test_wrapper, '', '',
+                                            argspec)
+    decorator = getattr(decorated, '_tf_decorator')
+    self.assertEqual(argspec, decorator.decorator_argspec)
+    self.assertEqual(
+        inspect.signature(decorated),
+        inspect.Signature([
+            inspect.Parameter('a', inspect.Parameter.POSITIONAL_OR_KEYWORD),
+            inspect.Parameter(
+                'b', inspect.Parameter.POSITIONAL_OR_KEYWORD, default=1
+            ),
+            inspect.Parameter(
+                'c',
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                default='hello',
+            ),
+            inspect.Parameter('args', inspect.Parameter.VAR_POSITIONAL),
+            inspect.Parameter('kwargs', inspect.Parameter.VAR_KEYWORD),
+        ]),
+    )
+
+  def testSetsDecoratorNameToFunctionThatCallsMakeDecoratorIfAbsent(self):
+
+    def test_decorator_name(wrapper):
+      return tf_decorator.make_decorator(test_function, wrapper)
+
+    decorated = test_decorator_name(test_wrapper)
+    decorator = getattr(decorated, '_tf_decorator')
+    self.assertEqual('test_decorator_name', decorator.decorator_name)
+
+  def testCompatibleWithNamelessCallables(self):
+
+    class Callable(object):
+
+      def __call__(self):
+        pass
+
+    callable_object = Callable()
+    # Smoke test: This should not raise an exception, even though
+    # `callable_object` does not have a `__name__` attribute.
+    _ = tf_decorator.make_decorator(callable_object, test_wrapper)
+
+    partial = functools.partial(test_function, x=1)
+    # Smoke test: This should not raise an exception, even though `partial` does
+    # not have `__name__`, `__module__`, and `__doc__` attributes.
+    _ = tf_decorator.make_decorator(partial, test_wrapper)
+
+
+class TfDecoratorRewrapTest(test.TestCase):
+
+  def testRewrapMutatesAffectedFunction(self):
+
+    @test_injectable_decorator_square
+    @test_injectable_decorator_increment
+    def test_rewrappable_decorated(x):
+      return x * 2
+
+    def new_target(x):
+      return x * 3
+
+    self.assertEqual((1 * 2 + 1)**2, test_rewrappable_decorated(1))
+    prev_target, _ = tf_decorator.unwrap(test_rewrappable_decorated)
+    tf_decorator.rewrap(test_rewrappable_decorated, prev_target, new_target)
+    self.assertEqual((1 * 3 + 1)**2, test_rewrappable_decorated(1))
+
+  def testRewrapOfDecoratorFunction(self):
+
+    @test_injectable_decorator_square
+    @test_injectable_decorator_increment
+    def test_rewrappable_decorated(x):
+      return x * 2
+
+    def new_target(x):
+      return x * 3
+
+    prev_target = test_rewrappable_decorated._tf_decorator._decorated_target
+    # In this case, only the outer decorator (test_injectable_decorator_square)
+    # should be preserved.
+    tf_decorator.rewrap(test_rewrappable_decorated, prev_target, new_target)
+    self.assertEqual((1 * 3)**2, test_rewrappable_decorated(1))
+
+
+class TfDecoratorUnwrapTest(test.TestCase):
+
+  def testUnwrapReturnsEmptyArrayForUndecoratedFunction(self):
+    decorators, _ = tf_decorator.unwrap(test_function)
+    self.assertEqual(0, len(decorators))
+
+  def testUnwrapReturnsUndecoratedFunctionAsTarget(self):
+    _, target = tf_decorator.unwrap(test_function)
+    self.assertIs(test_function, target)
+
+  def testUnwrapReturnsFinalFunctionAsTarget(self):
+    self.assertEqual((4 + 1) * 2, test_decorated_function(4))
+    _, target = tf_decorator.unwrap(test_decorated_function)
+    self.assertTrue(tf_inspect.isfunction(target))
+    self.assertEqual(4 * 2, target(4))
+
+  def testUnwrapReturnsListOfUniqueTFDecorators(self):
+    decorators, _ = tf_decorator.unwrap(test_decorated_function)
+    self.assertEqual(3, len(decorators))
+    self.assertTrue(isinstance(decorators[0], tf_decorator.TFDecorator))
+    self.assertTrue(isinstance(decorators[1], tf_decorator.TFDecorator))
+    self.assertTrue(isinstance(decorators[2], tf_decorator.TFDecorator))
+    self.assertIsNot(decorators[0], decorators[1])
+    self.assertIsNot(decorators[1], decorators[2])
+    self.assertIsNot(decorators[2], decorators[0])
+
+  def testUnwrapReturnsDecoratorListFromOutermostToInnermost(self):
+    decorators, _ = tf_decorator.unwrap(test_decorated_function)
+    self.assertEqual('decorator 1', decorators[0].decorator_name)
+    self.assertEqual('test_decorator_increment_first_int_arg',
+                     decorators[1].decorator_name)
+    self.assertEqual('decorator 3', decorators[2].decorator_name)
+    self.assertEqual('decorator 3 documentation', decorators[2].decorator_doc)
+
+  def testUnwrapBoundMethods(self):
+    test_decorated_class = TestDecoratedClass()
+    self.assertEqual([2, 2, 3], test_decorated_class.return_params(1, 2, 3))
+    decorators, target = tf_decorator.unwrap(test_decorated_class.return_params)
+    self.assertEqual('test_decorator_increment_first_int_arg',
+                     decorators[0].decorator_name)
+    self.assertEqual([1, 2, 3], target(test_decorated_class, 1, 2, 3))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/tf_export.py

@@ -0,0 +1,398 @@
+# 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 exporting TensorFlow symbols to the API.
+
+Exporting a function or a class:
+
+To export a function or a class use tf_export decorator. For e.g.:
+```python
+@tf_export('foo', 'bar.foo')
+def foo(...):
+  ...
+```
+
+If a function is assigned to a variable, you can export it by calling
+tf_export explicitly. For e.g.:
+```python
+foo = get_foo(...)
+tf_export('foo', 'bar.foo')(foo)
+```
+
+
+Exporting a constant
+```python
+foo = 1
+tf_export('consts.foo').export_constant(__name__, 'foo')
+```
+"""
+from collections.abc import Sequence
+import functools
+import sys
+from typing import Any, NamedTuple, Optional, Protocol, TypeVar
+
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+
+KERAS_API_NAME = 'keras'
+TENSORFLOW_API_NAME = 'tensorflow'
+
+# List of subpackage names used by TensorFlow components. Have to check that
+# TensorFlow core repo does not export any symbols under these names.
+SUBPACKAGE_NAMESPACES = []
+
+
+class _Attributes(NamedTuple):
+  names: str
+  constants: str
+
+
+# Attribute values must be unique to each API.
+API_ATTRS = {
+    TENSORFLOW_API_NAME: _Attributes('_tf_api_names', '_tf_api_constants'),
+    KERAS_API_NAME: _Attributes('_keras_api_names', '_keras_api_constants'),
+}
+
+API_ATTRS_V1 = {
+    TENSORFLOW_API_NAME: _Attributes(
+        '_tf_api_names_v1', '_tf_api_constants_v1'
+    ),
+    KERAS_API_NAME: _Attributes(
+        '_keras_api_names_v1', '_keras_api_constants_v1'
+    ),
+}
+
+
+class InvalidSymbolNameError(Exception):
+  """Raised when trying to export symbol as an invalid or unallowed name."""
+
+
+_NAME_TO_SYMBOL_MAPPING: dict[str, Any] = dict()
+
+
+def get_symbol_from_name(name: str) -> Optional[Any]:
+  return _NAME_TO_SYMBOL_MAPPING.get(name)
+
+
+def get_canonical_name_for_symbol(
+    symbol: Any,
+    api_name: str = TENSORFLOW_API_NAME,
+    add_prefix_to_v1_names: bool = False,
+) -> Optional[str]:
+  """Get canonical name for the API symbol.
+
+  Example:
+  ```python
+  from tensorflow.python.util import tf_export
+  cls = tf_export.get_symbol_from_name('keras.optimizers.Adam')
+
+  # Gives `<class 'keras.optimizer_v2.adam.Adam'>`
+  print(cls)
+
+  # Gives `keras.optimizers.Adam`
+  print(tf_export.get_canonical_name_for_symbol(cls, api_name='keras'))
+  ```
+
+  Args:
+    symbol: API function or class.
+    api_name: API name. Currently, only `tensorflow`.
+    add_prefix_to_v1_names: Specifies whether a name available only in V1 should
+      be prefixed with compat.v1.
+
+  Returns:
+    Canonical name for the API symbol (for e.g. initializers.zeros) if
+    canonical name could be determined. Otherwise, returns None.
+  """
+  if not hasattr(symbol, '__dict__'):
+    return None
+  api_names_attr = API_ATTRS[api_name].names
+  _, undecorated_symbol = tf_decorator.unwrap(symbol)
+  if api_names_attr not in undecorated_symbol.__dict__:
+    return None
+  api_names = getattr(undecorated_symbol, api_names_attr)
+  deprecated_api_names = undecorated_symbol.__dict__.get(
+      '_tf_deprecated_api_names', []
+  )
+
+  canonical_name = get_canonical_name(api_names, deprecated_api_names)
+  if canonical_name:
+    return canonical_name
+
+  # If there is no V2 canonical name, get V1 canonical name.
+  api_names_attr = API_ATTRS_V1[api_name].names
+  api_names = getattr(undecorated_symbol, api_names_attr)
+  v1_canonical_name = get_canonical_name(api_names, deprecated_api_names)
+  if add_prefix_to_v1_names:
+    return 'compat.v1.%s' % v1_canonical_name
+  return v1_canonical_name
+
+
+def get_canonical_name(
+    api_names: Sequence[str], deprecated_api_names: Sequence[str]
+) -> Optional[str]:
+  """Get preferred endpoint name.
+
+  Args:
+    api_names: API names iterable.
+    deprecated_api_names: Deprecated API names iterable.
+
+  Returns:
+    Returns one of the following in decreasing preference:
+    - first non-deprecated endpoint
+    - first endpoint
+    - None
+  """
+  non_deprecated_name = next(
+      (name for name in api_names if name not in deprecated_api_names), None
+  )
+  if non_deprecated_name:
+    return non_deprecated_name
+  if api_names:
+    return api_names[0]
+  return None
+
+
+def get_v1_names(symbol: Any) -> Sequence[str]:
+  """Get a list of TF 1.* names for this symbol.
+
+  Args:
+    symbol: symbol to get API names for.
+
+  Returns:
+    List of all API names for this symbol.
+  """
+  names_v1 = []
+  tensorflow_api_attr_v1 = API_ATTRS_V1[TENSORFLOW_API_NAME].names
+  keras_api_attr_v1 = API_ATTRS_V1[KERAS_API_NAME].names
+
+  if not hasattr(symbol, '__dict__'):
+    return names_v1
+  if tensorflow_api_attr_v1 in symbol.__dict__:
+    names_v1.extend(getattr(symbol, tensorflow_api_attr_v1))
+  if keras_api_attr_v1 in symbol.__dict__:
+    names_v1.extend(getattr(symbol, keras_api_attr_v1))
+  return names_v1
+
+
+def get_v2_names(symbol: Any) -> Sequence[str]:
+  """Get a list of TF 2.0 names for this symbol.
+
+  Args:
+    symbol: symbol to get API names for.
+
+  Returns:
+    List of all API names for this symbol.
+  """
+  names_v2 = []
+  tensorflow_api_attr = API_ATTRS[TENSORFLOW_API_NAME].names
+  keras_api_attr = API_ATTRS[KERAS_API_NAME].names
+
+  if not hasattr(symbol, '__dict__'):
+    return names_v2
+  if tensorflow_api_attr in symbol.__dict__:
+    names_v2.extend(getattr(symbol, tensorflow_api_attr))
+  if keras_api_attr in symbol.__dict__:
+    names_v2.extend(getattr(symbol, keras_api_attr))
+  return names_v2
+
+
+def get_v1_constants(module: Any) -> Sequence[str]:
+  """Get a list of TF 1.* constants in this module.
+
+  Args:
+    module: TensorFlow module.
+
+  Returns:
+    List of all API constants under the given module.
+  """
+  constants_v1 = []
+  tensorflow_constants_attr_v1 = API_ATTRS_V1[TENSORFLOW_API_NAME].constants
+
+  if hasattr(module, tensorflow_constants_attr_v1):
+    constants_v1.extend(getattr(module, tensorflow_constants_attr_v1))
+  return constants_v1
+
+
+def get_v2_constants(module: Any) -> Sequence[str]:
+  """Get a list of TF 2.0 constants in this module.
+
+  Args:
+    module: TensorFlow module.
+
+  Returns:
+    List of all API constants under the given module.
+  """
+  constants_v2 = []
+  tensorflow_constants_attr = API_ATTRS[TENSORFLOW_API_NAME].constants
+
+  if hasattr(module, tensorflow_constants_attr):
+    constants_v2.extend(getattr(module, tensorflow_constants_attr))
+  return constants_v2
+
+
+T = TypeVar('T')
+
+
+class api_export(object):  # pylint: disable=invalid-name
+  """Provides ways to export symbols to the TensorFlow API."""
+
+  _names: Sequence[str]
+  _names_v1: Sequence[str]
+  _api_name: str
+
+  def __init__(
+      self,
+      *args: str,
+      api_name: str = TENSORFLOW_API_NAME,
+      v1: Optional[Sequence[str]] = None,
+      allow_multiple_exports: bool = True,  # pylint: disable=unused-argument
+  ):
+    """Export under the names *args (first one is considered canonical).
+
+    Args:
+      *args: API names in dot delimited format.
+      api_name: API you want to generate Currently, only `tensorflow`.
+      v1: Names for the TensorFlow V1 API. If not set, we will use V2 API names
+        both for TensorFlow V1 and V2 APIs.
+      allow_multiple_exports: Deprecated.
+    """
+    self._names = args
+    self._names_v1 = v1 if v1 is not None else args
+    self._api_name = api_name
+
+    self._validate_symbol_names()
+
+  def _validate_symbol_names(self) -> None:
+    """Validate you are exporting symbols under an allowed package.
+
+    We need to ensure things exported by tf_export, etc.
+    export symbols under disjoint top-level package names.
+
+    For TensorFlow, we check that it does not export anything under subpackage
+    names used by components (keras, etc.).
+
+    For each component, we check that it exports everything under its own
+    subpackage.
+
+    Raises:
+      InvalidSymbolNameError: If you try to export symbol under disallowed name.
+    """
+    all_symbol_names = set(self._names) | set(self._names_v1)
+    if self._api_name == TENSORFLOW_API_NAME:
+      for subpackage in SUBPACKAGE_NAMESPACES:
+        if any(n.startswith(subpackage) for n in all_symbol_names):
+          raise InvalidSymbolNameError(
+              '@tf_export is not allowed to export symbols under %s.*'
+              % (subpackage)
+          )
+    else:
+      if not all(n.startswith(self._api_name) for n in all_symbol_names):
+        raise InvalidSymbolNameError(
+            'Can only export symbols under package name of component.'
+        )
+
+  def __call__(self, func: T) -> T:
+    """Calls this decorator.
+
+    Args:
+      func: decorated symbol (function or class).
+
+    Returns:
+      The input function with _tf_api_names attribute set.
+    """
+    api_names_attr = API_ATTRS[self._api_name].names
+    api_names_attr_v1 = API_ATTRS_V1[self._api_name].names
+
+    _, undecorated_func = tf_decorator.unwrap(func)
+    self.set_attr(undecorated_func, api_names_attr, self._names)
+    self.set_attr(undecorated_func, api_names_attr_v1, self._names_v1)
+
+    for name in self._names:
+      _NAME_TO_SYMBOL_MAPPING[name] = func
+    for name_v1 in self._names_v1:
+      _NAME_TO_SYMBOL_MAPPING['compat.v1.%s' % name_v1] = func
+
+    return func
+
+  def set_attr(
+      self, func: Any, api_names_attr: str, names: Sequence[str]
+  ) -> None:
+    setattr(func, api_names_attr, names)
+
+  def export_constant(self, module_name: str, name: str) -> None:
+    """Store export information for constants/string literals.
+
+    Export information is stored in the module where constants/string literals
+    are defined.
+
+    e.g.
+    ```python
+    foo = 1
+    bar = 2
+    tf_export("consts.foo").export_constant(__name__, 'foo')
+    tf_export("consts.bar").export_constant(__name__, 'bar')
+    ```
+
+    Args:
+      module_name: (string) Name of the module to store constant at.
+      name: (string) Current constant name.
+    """
+    module = sys.modules[module_name]
+    api_constants_attr = API_ATTRS[self._api_name].constants
+    api_constants_attr_v1 = API_ATTRS_V1[self._api_name].constants
+
+    if not hasattr(module, api_constants_attr):
+      setattr(module, api_constants_attr, [])
+    # pylint: disable=protected-access
+    getattr(module, api_constants_attr).append((self._names, name))
+
+    if not hasattr(module, api_constants_attr_v1):
+      setattr(module, api_constants_attr_v1, [])
+    getattr(module, api_constants_attr_v1).append((self._names_v1, name))
+
+
+def kwarg_only(f: Any) -> Any:
+  """A wrapper that throws away all non-kwarg arguments."""
+  f_argspec = tf_inspect.getfullargspec(f)
+
+  def wrapper(*args, **kwargs):
+    if args:
+      raise TypeError(
+          '{f} only takes keyword args (possible keys: {kwargs}). '
+          'Please pass these args as kwargs instead.'.format(
+              f=f.__name__, kwargs=f_argspec.args
+          )
+      )
+    return f(**kwargs)
+
+  return tf_decorator.make_decorator(f, wrapper, decorator_argspec=f_argspec)
+
+
+class ExportType(Protocol):
+
+  def __call__(
+      self,
+      *v2: str,
+      v1: Optional[Sequence[str]] = None,
+      allow_multiple_exports: bool = True,  # Deprecated, no-op
+  ) -> api_export:
+    ...
+
+
+tf_export: ExportType = functools.partial(
+    api_export, api_name=TENSORFLOW_API_NAME
+)
+keras_export: ExportType = functools.partial(
+    api_export, api_name=KERAS_API_NAME
+)

tensorflow/python/util/tf_export_test.py

@@ -0,0 +1,190 @@
+# 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.
+# ==============================================================================
+"""tf_export tests."""
+
+# pylint: disable=unused-import
+import sys
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_export
+
+
+class ValidateExportTest(test.TestCase):
+  """Tests for tf_export class."""
+
+  class MockModule(object):
+
+    def __init__(self, name):
+      self.__name__ = name
+
+  def setUp(self):
+    super().setUp()
+    self._modules = []
+
+    def _test_function(unused_arg=0):
+      pass
+
+    def _test_function2(unused_arg=0):
+      pass
+
+    class TestClassA(object):
+      pass
+
+    class TestClassB(TestClassA):
+      pass
+
+    self._test_function = _test_function
+    self._test_function2 = _test_function2
+    self._test_class_a = TestClassA
+    self._test_class_b = TestClassB
+
+  def tearDown(self):
+    super().tearDown()
+    for name in self._modules:
+      del sys.modules[name]
+    self._modules = []
+
+  def _CreateMockModule(self, name):
+    mock_module = self.MockModule(name)
+    sys.modules[name] = mock_module
+    self._modules.append(name)
+    return mock_module
+
+  def testExportSingleFunction(self):
+    export_decorator = tf_export.tf_export('nameA', 'nameB')
+    decorated_function = export_decorator(self._test_function)
+    self.assertEqual(decorated_function, self._test_function)
+    self.assertEqual(('nameA', 'nameB'), decorated_function._tf_api_names)
+    self.assertEqual(['nameA', 'nameB'],
+                     tf_export.get_v1_names(decorated_function))
+    self.assertEqual(['nameA', 'nameB'],
+                     tf_export.get_v2_names(decorated_function))
+    self.assertEqual(
+        tf_export.get_symbol_from_name('nameA'), decorated_function)
+    self.assertEqual(
+        tf_export.get_symbol_from_name('nameB'), decorated_function)
+    self.assertEqual(
+        tf_export.get_symbol_from_name(
+            tf_export.get_canonical_name_for_symbol(decorated_function)),
+        decorated_function)
+
+  def testExportSingleFunctionV1Only(self):
+    export_decorator = tf_export.tf_export(v1=['nameA', 'nameB'])
+    decorated_function = export_decorator(self._test_function)
+    self.assertEqual(decorated_function, self._test_function)
+    self.assertAllEqual(('nameA', 'nameB'), decorated_function._tf_api_names_v1)
+    self.assertAllEqual(['nameA', 'nameB'],
+                        tf_export.get_v1_names(decorated_function))
+    self.assertEqual([], tf_export.get_v2_names(decorated_function))
+    self.assertEqual(
+        tf_export.get_symbol_from_name('compat.v1.nameA'), decorated_function)
+    self.assertEqual(
+        tf_export.get_symbol_from_name('compat.v1.nameB'), decorated_function)
+    self.assertEqual(
+        tf_export.get_symbol_from_name(
+            tf_export.get_canonical_name_for_symbol(
+                decorated_function, add_prefix_to_v1_names=True)),
+        decorated_function)
+
+  def testExportMultipleFunctions(self):
+    export_decorator1 = tf_export.tf_export('nameA', 'nameB')
+    export_decorator2 = tf_export.tf_export('nameC', 'nameD')
+    decorated_function1 = export_decorator1(self._test_function)
+    decorated_function2 = export_decorator2(self._test_function2)
+    self.assertEqual(decorated_function1, self._test_function)
+    self.assertEqual(decorated_function2, self._test_function2)
+    self.assertEqual(('nameA', 'nameB'), decorated_function1._tf_api_names)
+    self.assertEqual(('nameC', 'nameD'), decorated_function2._tf_api_names)
+    self.assertEqual(
+        tf_export.get_symbol_from_name('nameB'), decorated_function1)
+    self.assertEqual(
+        tf_export.get_symbol_from_name('nameD'), decorated_function2)
+    self.assertEqual(
+        tf_export.get_symbol_from_name(
+            tf_export.get_canonical_name_for_symbol(decorated_function1)),
+        decorated_function1)
+    self.assertEqual(
+        tf_export.get_symbol_from_name(
+            tf_export.get_canonical_name_for_symbol(decorated_function2)),
+        decorated_function2)
+
+  def testExportClasses(self):
+    export_decorator_a = tf_export.tf_export('TestClassA1')
+    export_decorator_a(self._test_class_a)
+    self.assertEqual(('TestClassA1',), self._test_class_a._tf_api_names)
+    self.assertNotIn('_tf_api_names', self._test_class_b.__dict__)
+
+    export_decorator_b = tf_export.tf_export('TestClassB1')
+    export_decorator_b(self._test_class_b)
+    self.assertEqual(('TestClassA1',), self._test_class_a._tf_api_names)
+    self.assertEqual(('TestClassB1',), self._test_class_b._tf_api_names)
+    self.assertEqual(
+        ['TestClassA1'], tf_export.get_v1_names(self._test_class_a)
+    )
+    self.assertEqual(
+        ['TestClassB1'], tf_export.get_v1_names(self._test_class_b)
+    )
+
+  def testExportSingleConstant(self):
+    module1 = self._CreateMockModule('module1')
+
+    export_decorator = tf_export.tf_export('NAME_A', 'NAME_B')
+    export_decorator.export_constant('module1', 'test_constant')
+    self.assertEqual([(('NAME_A', 'NAME_B'), 'test_constant')],
+                     module1._tf_api_constants)
+    self.assertEqual([(('NAME_A', 'NAME_B'), 'test_constant')],
+                     tf_export.get_v1_constants(module1))
+    self.assertEqual([(('NAME_A', 'NAME_B'), 'test_constant')],
+                     tf_export.get_v2_constants(module1))
+
+  def testExportMultipleConstants(self):
+    module1 = self._CreateMockModule('module1')
+    module2 = self._CreateMockModule('module2')
+
+    test_constant1 = 123
+    test_constant2 = 'abc'
+    test_constant3 = 0.5
+
+    export_decorator1 = tf_export.tf_export('NAME_A', 'NAME_B')
+    export_decorator2 = tf_export.tf_export('NAME_C', 'NAME_D')
+    export_decorator3 = tf_export.tf_export('NAME_E', 'NAME_F')
+    export_decorator1.export_constant('module1', test_constant1)
+    export_decorator2.export_constant('module2', test_constant2)
+    export_decorator3.export_constant('module2', test_constant3)
+    self.assertEqual([(('NAME_A', 'NAME_B'), 123)], module1._tf_api_constants)
+    self.assertEqual([(('NAME_C', 'NAME_D'), 'abc'),
+                      (('NAME_E', 'NAME_F'), 0.5)], module2._tf_api_constants)
+
+  def testMultipleDecorators(self):
+
+    def get_wrapper(func):
+
+      def wrapper(*unused_args, **unused_kwargs):
+        pass
+
+      return tf_decorator.make_decorator(func, wrapper)
+
+    decorated_function = get_wrapper(self._test_function)
+
+    export_decorator = tf_export.tf_export('nameA', 'nameB')
+    exported_function = export_decorator(decorated_function)
+    self.assertEqual(decorated_function, exported_function)
+    self.assertEqual(('nameA', 'nameB'), self._test_function._tf_api_names)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/tf_inspect.py

@@ -0,0 +1,470 @@
+# 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.
+# ==============================================================================
+"""TFDecorator-aware replacements for the inspect module."""
+import collections
+import functools
+import inspect as _inspect
+
+from tensorflow.python.util import tf_decorator
+
+
+# inspect.signature() is preferred over inspect.getfullargspec() in PY3.
+# Note that while it can handle TFDecorators, it will ignore a TFDecorator's
+# provided ArgSpec/FullArgSpec and instead return the signature of the
+# inner-most function.
+def signature(obj, *, follow_wrapped=True):
+  """TFDecorator-aware replacement for inspect.signature."""
+  return _inspect.signature(
+      tf_decorator.unwrap(obj)[1], follow_wrapped=follow_wrapped)
+
+
+Parameter = _inspect.Parameter
+Signature = _inspect.Signature
+
+if hasattr(_inspect, 'ArgSpec'):
+  ArgSpec = _inspect.ArgSpec
+else:
+  ArgSpec = collections.namedtuple(
+      'ArgSpec',
+      [
+          'args',
+          'varargs',
+          'keywords',
+          'defaults',
+      ],
+  )
+
+
+if hasattr(_inspect, 'FullArgSpec'):
+  FullArgSpec = _inspect.FullArgSpec  # pylint: disable=invalid-name
+else:
+  FullArgSpec = collections.namedtuple('FullArgSpec', [
+      'args', 'varargs', 'varkw', 'defaults', 'kwonlyargs', 'kwonlydefaults',
+      'annotations'
+  ])
+
+
+def _convert_maybe_argspec_to_fullargspec(argspec):
+  if isinstance(argspec, FullArgSpec):
+    return argspec
+  return FullArgSpec(
+      args=argspec.args,
+      varargs=argspec.varargs,
+      varkw=argspec.keywords,
+      defaults=argspec.defaults,
+      kwonlyargs=[],
+      kwonlydefaults=None,
+      annotations={})
+
+if hasattr(_inspect, 'getfullargspec'):
+  _getfullargspec = _inspect.getfullargspec  # pylint: disable=invalid-name
+
+  def _getargspec(target):
+    """A python3 version of getargspec.
+
+    Calls `getfullargspec` and assigns args, varargs,
+    varkw, and defaults to a python 2/3 compatible `ArgSpec`.
+
+    The parameter name 'varkw' is changed to 'keywords' to fit the
+    `ArgSpec` struct.
+
+    Args:
+      target: the target object to inspect.
+
+    Returns:
+      An ArgSpec with args, varargs, keywords, and defaults parameters
+      from FullArgSpec.
+    """
+    fullargspecs = getfullargspec(target)
+
+    defaults = fullargspecs.defaults or ()
+    if fullargspecs.kwonlydefaults:
+      defaults += tuple(fullargspecs.kwonlydefaults.values())
+
+    if not defaults:
+      defaults = None
+
+    argspecs = ArgSpec(
+        args=fullargspecs.args + fullargspecs.kwonlyargs,
+        varargs=fullargspecs.varargs,
+        keywords=fullargspecs.varkw,
+        defaults=defaults,
+    )
+    return argspecs
+else:
+  _getargspec = _inspect.getargspec
+
+  def _getfullargspec(target):
+    """A python2 version of getfullargspec.
+
+    Args:
+      target: the target object to inspect.
+
+    Returns:
+      A FullArgSpec with empty kwonlyargs, kwonlydefaults and annotations.
+    """
+    return _convert_maybe_argspec_to_fullargspec(getargspec(target))
+
+
+def currentframe():
+  """TFDecorator-aware replacement for inspect.currentframe."""
+  return _inspect.stack()[1][0]
+
+
+def getargspec(obj):
+  """TFDecorator-aware replacement for `inspect.getargspec`.
+
+  Note: `getfullargspec` is recommended as the python 2/3 compatible
+  replacement for this function.
+
+  Args:
+    obj: A function, partial function, or callable object, possibly decorated.
+
+  Returns:
+    The `ArgSpec` that describes the signature of the outermost decorator that
+    changes the callable's signature, or the `ArgSpec` that describes
+    the object if not decorated.
+
+  Raises:
+    ValueError: When callable's signature can not be expressed with
+      ArgSpec.
+    TypeError: For objects of unsupported types.
+  """
+  if isinstance(obj, functools.partial):
+    return _get_argspec_for_partial(obj)
+
+  decorators, target = tf_decorator.unwrap(obj)
+
+  spec = next((d.decorator_argspec
+               for d in decorators
+               if d.decorator_argspec is not None), None)
+  if spec:
+    return spec
+
+  try:
+    # Python3 will handle most callables here (not partial).
+    return _getargspec(target)
+  except TypeError:
+    pass
+
+  if isinstance(target, type):
+    try:
+      return _getargspec(target.__init__)
+    except TypeError:
+      pass
+
+    try:
+      return _getargspec(target.__new__)
+    except TypeError:
+      pass
+
+  # The `type(target)` ensures that if a class is received we don't return
+  # the signature of its __call__ method.
+  return _getargspec(type(target).__call__)
+
+
+def _get_argspec_for_partial(obj):
+  """Implements `getargspec` for `functools.partial` objects.
+
+  Args:
+    obj: The `functools.partial` object
+  Returns:
+    An `inspect.ArgSpec`
+  Raises:
+    ValueError: When callable's signature can not be expressed with
+      ArgSpec.
+  """
+  # When callable is a functools.partial object, we construct its ArgSpec with
+  # following strategy:
+  # - If callable partial contains default value for positional arguments (ie.
+  # object.args), then final ArgSpec doesn't contain those positional arguments.
+  # - If callable partial contains default value for keyword arguments (ie.
+  # object.keywords), then we merge them with wrapped target. Default values
+  # from callable partial takes precedence over those from wrapped target.
+  #
+  # However, there is a case where it is impossible to construct a valid
+  # ArgSpec. Python requires arguments that have no default values must be
+  # defined before those with default values. ArgSpec structure is only valid
+  # when this presumption holds true because default values are expressed as a
+  # tuple of values without keywords and they are always assumed to belong to
+  # last K arguments where K is number of default values present.
+  #
+  # Since functools.partial can give default value to any argument, this
+  # presumption may no longer hold in some cases. For example:
+  #
+  # def func(m, n):
+  #   return 2 * m + n
+  # partialed = functools.partial(func, m=1)
+  #
+  # This example will result in m having a default value but n doesn't. This is
+  # usually not allowed in Python and can not be expressed in ArgSpec correctly.
+  #
+  # Thus, we must detect cases like this by finding first argument with default
+  # value and ensures all following arguments also have default values. When
+  # this is not true, a ValueError is raised.
+
+  n_prune_args = len(obj.args)
+  partial_keywords = obj.keywords or {}
+
+  args, varargs, keywords, defaults = getargspec(obj.func)
+
+  # Pruning first n_prune_args arguments.
+  args = args[n_prune_args:]
+
+  # Partial function may give default value to any argument, therefore length
+  # of default value list must be len(args) to allow each argument to
+  # potentially be given a default value.
+  no_default = object()
+  all_defaults = [no_default] * len(args)
+
+  if defaults:
+    all_defaults[-len(defaults):] = defaults
+
+  # Fill in default values provided by partial function in all_defaults.
+  for kw, default in iter(partial_keywords.items()):
+    if kw in args:
+      idx = args.index(kw)
+      all_defaults[idx] = default
+    elif not keywords:
+      raise ValueError(f'{obj} does not have a **kwargs parameter, but '
+                       f'contains an unknown partial keyword {kw}.')
+
+  # Find first argument with default value set.
+  first_default = next(
+      (idx for idx, x in enumerate(all_defaults) if x is not no_default), None)
+
+  # If no default values are found, return ArgSpec with defaults=None.
+  if first_default is None:
+    return ArgSpec(args, varargs, keywords, None)
+
+  # Checks if all arguments have default value set after first one.
+  invalid_default_values = [
+      args[i] for i, j in enumerate(all_defaults)
+      if j is no_default and i > first_default
+  ]
+
+  if invalid_default_values:
+    raise ValueError(f'{obj} has some keyword-only arguments, which are not'
+                     f' supported: {invalid_default_values}.')
+
+  return ArgSpec(args, varargs, keywords, tuple(all_defaults[first_default:]))
+
+
+def getfullargspec(obj):
+  """TFDecorator-aware replacement for `inspect.getfullargspec`.
+
+  This wrapper emulates `inspect.getfullargspec` in[^)]* Python2.
+
+  Args:
+    obj: A callable, possibly decorated.
+
+  Returns:
+    The `FullArgSpec` that describes the signature of
+    the outermost decorator that changes the callable's signature. If the
+    callable is not decorated, `inspect.getfullargspec()` will be called
+    directly on the callable.
+  """
+  decorators, target = tf_decorator.unwrap(obj)
+
+  for d in decorators:
+    if d.decorator_argspec is not None:
+      return _convert_maybe_argspec_to_fullargspec(d.decorator_argspec)
+  return _getfullargspec(target)
+
+
+def getcallargs(*func_and_positional, **named):
+  """TFDecorator-aware replacement for inspect.getcallargs.
+
+  Args:
+    *func_and_positional: A callable, possibly decorated, followed by any
+      positional arguments that would be passed to `func`.
+    **named: The named argument dictionary that would be passed to `func`.
+
+  Returns:
+    A dictionary mapping `func`'s named arguments to the values they would
+    receive if `func(*positional, **named)` were called.
+
+  `getcallargs` will use the argspec from the outermost decorator that provides
+  it. If no attached decorators modify argspec, the final unwrapped target's
+  argspec will be used.
+  """
+  func = func_and_positional[0]
+  positional = func_and_positional[1:]
+  argspec = getfullargspec(func)
+  call_args = named.copy()
+  this = getattr(func, 'im_self', None) or getattr(func, '__self__', None)
+  if ismethod(func) and this:
+    positional = (this,) + positional
+  remaining_positionals = [arg for arg in argspec.args if arg not in call_args]
+  call_args.update(dict(zip(remaining_positionals, positional)))
+  default_count = 0 if not argspec.defaults else len(argspec.defaults)
+  if default_count:
+    for arg, value in zip(argspec.args[-default_count:], argspec.defaults):
+      if arg not in call_args:
+        call_args[arg] = value
+  if argspec.kwonlydefaults is not None:
+    for k, v in argspec.kwonlydefaults.items():
+      if k not in call_args:
+        call_args[k] = v
+  return call_args
+
+
+def getframeinfo(*args, **kwargs):
+  return _inspect.getframeinfo(*args, **kwargs)
+
+
+def getdoc(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getdoc.
+
+  Args:
+    object: An object, possibly decorated.
+
+  Returns:
+    The docstring associated with the object.
+
+  The outermost-decorated object is intended to have the most complete
+  documentation, so the decorated parameter is not unwrapped.
+  """
+  return _inspect.getdoc(object)
+
+
+def getfile(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getfile."""
+  unwrapped_object = tf_decorator.unwrap(object)[1]
+
+  # Work around for the case when object is a stack frame
+  # and only .pyc files are used. In this case, getfile
+  # might return incorrect path. So, we get the path from f_globals
+  # instead.
+  if (hasattr(unwrapped_object, 'f_globals') and
+      '__file__' in unwrapped_object.f_globals):
+    return unwrapped_object.f_globals['__file__']
+  return _inspect.getfile(unwrapped_object)
+
+
+def getmembers(object, predicate=None):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getmembers."""
+  return _inspect.getmembers(object, predicate)
+
+
+def getmodule(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getmodule."""
+  return _inspect.getmodule(object)
+
+
+def getmro(cls):
+  """TFDecorator-aware replacement for inspect.getmro."""
+  return _inspect.getmro(cls)
+
+
+def getsource(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getsource."""
+  return _inspect.getsource(tf_decorator.unwrap(object)[1])
+
+
+def getsourcefile(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getsourcefile."""
+  return _inspect.getsourcefile(tf_decorator.unwrap(object)[1])
+
+
+def getsourcelines(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.getsourcelines."""
+  return _inspect.getsourcelines(tf_decorator.unwrap(object)[1])
+
+
+def isbuiltin(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.isbuiltin."""
+  return _inspect.isbuiltin(tf_decorator.unwrap(object)[1])
+
+
+def isclass(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.isclass."""
+  return _inspect.isclass(tf_decorator.unwrap(object)[1])
+
+
+def isfunction(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.isfunction."""
+  return _inspect.isfunction(tf_decorator.unwrap(object)[1])
+
+
+def isframe(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.ismodule."""
+  return _inspect.isframe(tf_decorator.unwrap(object)[1])
+
+
+def isgenerator(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.isgenerator."""
+  return _inspect.isgenerator(tf_decorator.unwrap(object)[1])
+
+
+def isgeneratorfunction(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.isgeneratorfunction."""
+  return _inspect.isgeneratorfunction(tf_decorator.unwrap(object)[1])
+
+
+def ismethod(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.ismethod."""
+  return _inspect.ismethod(tf_decorator.unwrap(object)[1])
+
+
+def isanytargetmethod(object):  # pylint: disable=redefined-builtin
+  # pylint: disable=g-doc-args,g-doc-return-or-yield
+  """Checks if `object` or a TF Decorator wrapped target contains self or cls.
+
+  This function could be used along with `tf_inspect.getfullargspec` to
+  determine if the first argument of `object` argspec is self or cls. If the
+  first argument is self or cls, it needs to be excluded from argspec when we
+  compare the argspec to the input arguments and, if provided, the tf.function
+  input_signature.
+
+  Like `tf_inspect.getfullargspec` and python `inspect.getfullargspec`, it
+  does not unwrap python decorators.
+
+  Args:
+    obj: An method, function, or functool.partial, possibly decorated by
+    TFDecorator.
+
+  Returns:
+    A bool indicates if `object` or any target along the chain of TF decorators
+    is a method.
+  """
+  decorators, target = tf_decorator.unwrap(object)
+  for decorator in decorators:
+    if _inspect.ismethod(decorator.decorated_target):
+      return True
+
+  # TODO(b/194845243): Implement the long term solution with inspect.signature.
+  # A functools.partial object is not a function or method. But if the wrapped
+  # func is a method, the argspec will contain self/cls.
+  while isinstance(target, functools.partial):
+    target = target.func
+
+  # `target` is a method or an instance with __call__
+  return callable(target) and not _inspect.isfunction(target)
+
+
+def ismodule(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.ismodule."""
+  return _inspect.ismodule(tf_decorator.unwrap(object)[1])
+
+
+def isroutine(object):  # pylint: disable=redefined-builtin
+  """TFDecorator-aware replacement for inspect.isroutine."""
+  return _inspect.isroutine(tf_decorator.unwrap(object)[1])
+
+
+def stack(context=1):
+  """TFDecorator-aware replacement for inspect.stack."""
+  return _inspect.stack(context)[1:]

tensorflow/python/util/tf_inspect_test.py

@@ -0,0 +1,860 @@
+# 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.
+# ==============================================================================
+"""Unit tests for tf_inspect."""
+
+# pylint: disable=unused-import
+import functools
+import inspect
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+
+
+def test_decorator(decorator_name, decorator_doc=None):
+
+  def make_tf_decorator(target):
+    return tf_decorator.TFDecorator(decorator_name, target, decorator_doc)
+
+  return make_tf_decorator
+
+
+def test_undecorated_function():
+  pass
+
+
+@test_decorator('decorator 1')
+@test_decorator('decorator 2')
+@test_decorator('decorator 3')
+def test_decorated_function(x):
+  """Test Decorated Function Docstring."""
+  return x * 2
+
+
+@test_decorator('decorator')
+def test_decorated_function_with_defaults(a, b=2, c='Hello'):
+  """Test Decorated Function With Defaults Docstring."""
+  return [a, b, c]
+
+
+@test_decorator('decorator')
+def test_decorated_function_with_varargs_and_kwonlyargs(*args, b=2, c='Hello'):
+  """Test Decorated Function With both varargs and keyword args."""
+  return [args, b, c]
+
+
+@test_decorator('decorator')
+class TestDecoratedClass(object):
+  """Test Decorated Class."""
+
+  def __init__(self):
+    pass
+
+  def two(self):
+    return 2
+
+
+class TfInspectTest(test.TestCase):
+
+  def testCurrentFrame(self):
+    self.assertEqual(inspect.currentframe(), tf_inspect.currentframe())
+
+  def testGetArgSpecOnDecoratorsThatDontProvideArgspec(self):
+    argspec = tf_inspect.getargspec(test_decorated_function_with_defaults)
+    self.assertEqual(['a', 'b', 'c'], argspec.args)
+    self.assertEqual((2, 'Hello'), argspec.defaults)
+
+  def testGetArgSpecOnDecoratorThatChangesArgspec(self):
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={},
+    )
+
+    decorator = tf_decorator.TFDecorator('', test_undecorated_function, '',
+                                         argspec)
+    self.assertEqual(argspec, tf_inspect.getfullargspec(decorator))
+
+  def testGetArgSpecIgnoresDecoratorsThatDontProvideArgspec(self):
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={},
+    )
+
+    inner_decorator = tf_decorator.TFDecorator(
+        '', test_undecorated_function, '', argspec
+    )
+    outer_decorator = tf_decorator.TFDecorator('', inner_decorator)
+    self.assertEqual(argspec, tf_inspect.getargspec(outer_decorator))
+
+  def testGetArgSpecThatContainsVarargsAndKwonlyArgs(self):
+    argspec = tf_inspect.getargspec(
+        test_decorated_function_with_varargs_and_kwonlyargs
+    )
+    self.assertEqual(['b', 'c'], argspec.args)
+    self.assertEqual((2, 'Hello'), argspec.defaults)
+
+  def testGetArgSpecReturnsOutermostDecoratorThatChangesArgspec(self):
+    outer_argspec = tf_inspect.FullArgSpec(
+        args=['a'],
+        varargs=None,
+        varkw=None,
+        defaults=(),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={},
+    )
+    inner_argspec = tf_inspect.FullArgSpec(
+        args=['b'],
+        varargs=None,
+        varkw=None,
+        defaults=(),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={},
+    )
+
+    inner_decorator = tf_decorator.TFDecorator('', test_undecorated_function,
+                                               '', inner_argspec)
+    outer_decorator = tf_decorator.TFDecorator('', inner_decorator, '',
+                                               outer_argspec)
+    self.assertEqual(outer_argspec, tf_inspect.getfullargspec(outer_decorator))
+
+  def testGetArgSpecOnPartialPositionalArgumentOnly(self):
+    """Tests getargspec on partial function with only positional arguments."""
+
+    def func(m, n):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, 7)
+    argspec = tf_inspect.ArgSpec(
+        args=['n'], varargs=None, keywords=None, defaults=None)
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialArgumentWithConvertibleToFalse(self):
+    """Tests getargspec on partial function with args that convert to False."""
+
+    def func(m, n):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, m=0)
+
+    with self.assertRaisesRegex(ValueError, 'keyword-only arguments'):
+      tf_inspect.getargspec(partial_func)
+
+  def testGetArgSpecOnPartialInvalidArgspec(self):
+    """Tests getargspec on partial function that doesn't have valid argspec."""
+
+    def func(m, n, l, k=4):
+      return 2 * m + l + n * k
+
+    partial_func = functools.partial(func, n=7)
+
+    with self.assertRaisesRegex(ValueError, 'keyword-only arguments'):
+      tf_inspect.getargspec(partial_func)
+
+  def testGetArgSpecOnPartialValidArgspec(self):
+    """Tests getargspec on partial function with valid argspec."""
+
+    def func(m, n, l, k=4):
+      return 2 * m + l + n * k
+
+    partial_func = functools.partial(func, n=7, l=2)
+    argspec = tf_inspect.ArgSpec(
+        args=['m', 'n', 'l', 'k'],
+        varargs=None,
+        keywords=None,
+        defaults=(7, 2, 4))
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialNoArgumentsLeft(self):
+    """Tests getargspec on partial function that prunes all arguments."""
+
+    def func(m, n):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, 7, 10)
+    argspec = tf_inspect.ArgSpec(
+        args=[], varargs=None, keywords=None, defaults=None)
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialKeywordArgument(self):
+    """Tests getargspec on partial function that prunes some arguments."""
+
+    def func(m, n):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, n=7)
+    argspec = tf_inspect.ArgSpec(
+        args=['m', 'n'], varargs=None, keywords=None, defaults=(7,))
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialKeywordArgumentWithDefaultValue(self):
+    """Tests getargspec on partial function that prunes argument by keyword."""
+
+    def func(m=1, n=2):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, n=7)
+    argspec = tf_inspect.ArgSpec(
+        args=['m', 'n'], varargs=None, keywords=None, defaults=(1, 7))
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialWithVarargs(self):
+    """Tests getargspec on partial function with variable arguments."""
+
+    def func(m, *arg):
+      return m + len(arg)
+
+    partial_func = functools.partial(func, 7, 8)
+    argspec = tf_inspect.ArgSpec(
+        args=[], varargs='arg', keywords=None, defaults=None)
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialWithVarkwargs(self):
+    """Tests getargspec on partial function with variable keyword arguments."""
+
+    def func(m, n, **kwarg):
+      return m * n + len(kwarg)
+
+    partial_func = functools.partial(func, 7)
+    argspec = tf_inspect.ArgSpec(
+        args=['n'], varargs=None, keywords='kwarg', defaults=None)
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialWithDecorator(self):
+    """Tests getargspec on decorated partial function."""
+
+    @test_decorator('decorator')
+    def func(m=1, n=2):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, n=7)
+    argspec = tf_inspect.ArgSpec(
+        args=['m', 'n'], varargs=None, keywords=None, defaults=(1, 7))
+
+    self.assertEqual(argspec, tf_inspect.getargspec(partial_func))
+
+  def testGetArgSpecOnPartialWithDecoratorThatChangesArgspec(self):
+    """Tests getargspec on partial function with decorated argspec."""
+
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={},
+    )
+    decorator = tf_decorator.TFDecorator('', test_undecorated_function, '',
+                                         argspec)
+    signature = inspect.Signature([
+        inspect.Parameter(
+            'a', inspect.Parameter.KEYWORD_ONLY, default=2
+        ),
+        inspect.Parameter(
+            'b', inspect.Parameter.KEYWORD_ONLY, default=1
+        ),
+        inspect.Parameter(
+            'c', inspect.Parameter.KEYWORD_ONLY, default='hello'
+        ),
+    ])
+    partial_with_decorator = functools.partial(decorator, a=2)
+    self.assertEqual(argspec, tf_inspect.getfullargspec(decorator))
+    self.assertEqual(signature, inspect.signature(partial_with_decorator))
+
+  def testGetArgSpecOnCallableObject(self):
+
+    class Callable(object):
+
+      def __call__(self, a, b=1, c='hello'):
+        pass
+
+    argspec = tf_inspect.ArgSpec(
+        args=['self', 'a', 'b', 'c'],
+        varargs=None,
+        keywords=None,
+        defaults=(1, 'hello'))
+
+    test_obj = Callable()
+    self.assertEqual(argspec, tf_inspect.getargspec(test_obj))
+
+  def testGetArgSpecOnInitClass(self):
+
+    class InitClass(object):
+
+      def __init__(self, a, b=1, c='hello'):
+        pass
+
+    argspec = tf_inspect.ArgSpec(
+        args=['self', 'a', 'b', 'c'],
+        varargs=None,
+        keywords=None,
+        defaults=(1, 'hello'))
+
+    self.assertEqual(argspec, tf_inspect.getargspec(InitClass))
+
+  def testGetArgSpecOnNewClass(self):
+
+    class NewClass(object):
+
+      def __new__(cls, a, b=1, c='hello'):
+        pass
+
+    argspec = tf_inspect.ArgSpec(
+        args=['cls', 'a', 'b', 'c'],
+        varargs=None,
+        keywords=None,
+        defaults=(1, 'hello'))
+
+    self.assertEqual(argspec, tf_inspect.getargspec(NewClass))
+
+  def testGetFullArgSpecOnDecoratorsThatDontProvideFullArgSpec(self):
+    argspec = tf_inspect.getfullargspec(test_decorated_function_with_defaults)
+    self.assertEqual(['a', 'b', 'c'], argspec.args)
+    self.assertEqual((2, 'Hello'), argspec.defaults)
+
+  def testGetFullArgSpecOnDecoratorThatChangesFullArgSpec(self):
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    decorator = tf_decorator.TFDecorator('', test_undecorated_function, '',
+                                         argspec)
+    self.assertEqual(argspec, tf_inspect.getfullargspec(decorator))
+
+  def testGetFullArgSpecIgnoresDecoratorsThatDontProvideFullArgSpec(self):
+    argspec = tf_inspect.FullArgSpec(
+        args=['a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    inner_decorator = tf_decorator.TFDecorator('', test_undecorated_function,
+                                               '', argspec)
+    outer_decorator = tf_decorator.TFDecorator('', inner_decorator)
+    self.assertEqual(argspec, tf_inspect.getfullargspec(outer_decorator))
+
+  def testGetFullArgSpecReturnsOutermostDecoratorThatChangesFullArgSpec(self):
+    outer_argspec = tf_inspect.FullArgSpec(
+        args=['a'],
+        varargs=None,
+        varkw=None,
+        defaults=None,
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+    inner_argspec = tf_inspect.FullArgSpec(
+        args=['b'],
+        varargs=None,
+        varkw=None,
+        defaults=None,
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    inner_decorator = tf_decorator.TFDecorator('', test_undecorated_function,
+                                               '', inner_argspec)
+    outer_decorator = tf_decorator.TFDecorator('', inner_decorator, '',
+                                               outer_argspec)
+    self.assertEqual(outer_argspec, tf_inspect.getfullargspec(outer_decorator))
+
+  def testGetFullArgsSpecForPartial(self):
+
+    def func(a, b):
+      del a, b
+
+    partial_function = functools.partial(func, 1)
+    argspec = tf_inspect.FullArgSpec(
+        args=['b'],
+        varargs=None,
+        varkw=None,
+        defaults=None,
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    self.assertEqual(argspec, tf_inspect.getfullargspec(partial_function))
+
+  def testGetFullArgSpecOnPartialNoArgumentsLeft(self):
+    """Tests getfullargspec on partial function that prunes all arguments."""
+
+    def func(m, n):
+      return 2 * m + n
+
+    partial_func = functools.partial(func, 7, 10)
+    argspec = tf_inspect.FullArgSpec(
+        args=[],
+        varargs=None,
+        varkw=None,
+        defaults=None,
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    self.assertEqual(argspec, tf_inspect.getfullargspec(partial_func))
+
+  def testGetFullArgSpecOnPartialWithVarargs(self):
+    """Tests getfullargspec on partial function with variable arguments."""
+
+    def func(m, *arg):
+      return m + len(arg)
+
+    partial_func = functools.partial(func, 7, 8)
+    argspec = tf_inspect.FullArgSpec(
+        args=[],
+        varargs='arg',
+        varkw=None,
+        defaults=None,
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    self.assertEqual(argspec, tf_inspect.getfullargspec(partial_func))
+
+  def testGetFullArgSpecOnPartialWithVarkwargs(self):
+    """Tests getfullargspec.
+
+    Tests on partial function with variable keyword arguments.
+    """
+
+    def func(m, n, **kwarg):
+      return m * n + len(kwarg)
+
+    partial_func = functools.partial(func, 7)
+    argspec = tf_inspect.FullArgSpec(
+        args=['n'],
+        varargs=None,
+        varkw='kwarg',
+        defaults=None,
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    self.assertEqual(argspec, tf_inspect.getfullargspec(partial_func))
+
+  def testGetFullArgSpecOnCallableObject(self):
+
+    class Callable(object):
+
+      def __call__(self, a, b=1, c='hello'):
+        pass
+
+    argspec = tf_inspect.FullArgSpec(
+        args=['self', 'a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    test_obj = Callable()
+    self.assertEqual(argspec, tf_inspect.getfullargspec(test_obj))
+
+  def testGetFullArgSpecOnInitClass(self):
+
+    class InitClass(object):
+
+      def __init__(self, a, b=1, c='hello'):
+        pass
+
+    argspec = tf_inspect.FullArgSpec(
+        args=['self', 'a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    self.assertEqual(argspec, tf_inspect.getfullargspec(InitClass))
+
+  def testGetFullArgSpecOnNewClass(self):
+
+    class NewClass(object):
+
+      def __new__(cls, a, b=1, c='hello'):
+        pass
+
+    argspec = tf_inspect.FullArgSpec(
+        args=['cls', 'a', 'b', 'c'],
+        varargs=None,
+        varkw=None,
+        defaults=(1, 'hello'),
+        kwonlyargs=[],
+        kwonlydefaults=None,
+        annotations={})
+
+    self.assertEqual(argspec, tf_inspect.getfullargspec(NewClass))
+
+  def testSignatureOnDecoratorsThatDontProvideFullArgSpec(self):
+    signature = tf_inspect.signature(test_decorated_function_with_defaults)
+
+    self.assertEqual([
+        tf_inspect.Parameter('a', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        tf_inspect.Parameter(
+            'b', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD, default=2),
+        tf_inspect.Parameter(
+            'c', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD, default='Hello')
+    ], list(signature.parameters.values()))
+
+  def testSignatureFollowsNestedDecorators(self):
+    signature = tf_inspect.signature(test_decorated_function)
+
+    self.assertEqual(
+        [tf_inspect.Parameter('x', tf_inspect.Parameter.POSITIONAL_OR_KEYWORD)],
+        list(signature.parameters.values()))
+
+  def testGetDoc(self):
+    self.assertEqual('Test Decorated Function With Defaults Docstring.',
+                     tf_inspect.getdoc(test_decorated_function_with_defaults))
+
+  def testGetFile(self):
+    self.assertTrue('tf_inspect_test.py' in tf_inspect.getfile(
+        test_decorated_function_with_defaults))
+    self.assertTrue('tf_decorator.py' in tf_inspect.getfile(
+        test_decorator('decorator')(tf_decorator.unwrap)))
+
+  def testGetMembers(self):
+    self.assertEqual(
+        inspect.getmembers(TestDecoratedClass),
+        tf_inspect.getmembers(TestDecoratedClass))
+
+  def testGetModule(self):
+    self.assertEqual(
+        inspect.getmodule(TestDecoratedClass),
+        tf_inspect.getmodule(TestDecoratedClass))
+    self.assertEqual(
+        inspect.getmodule(test_decorated_function),
+        tf_inspect.getmodule(test_decorated_function))
+    self.assertEqual(
+        inspect.getmodule(test_undecorated_function),
+        tf_inspect.getmodule(test_undecorated_function))
+
+  def testGetSource(self):
+    expected = '''@test_decorator('decorator')
+def test_decorated_function_with_defaults(a, b=2, c='Hello'):
+  """Test Decorated Function With Defaults Docstring."""
+  return [a, b, c]
+'''
+    self.assertEqual(
+        expected, tf_inspect.getsource(test_decorated_function_with_defaults))
+
+  def testGetSourceFile(self):
+    self.assertEqual(
+        __file__,
+        tf_inspect.getsourcefile(test_decorated_function_with_defaults))
+
+  def testGetSourceLines(self):
+    expected = inspect.getsourcelines(
+        test_decorated_function_with_defaults.decorated_target)
+    self.assertEqual(
+        expected,
+        tf_inspect.getsourcelines(test_decorated_function_with_defaults))
+
+  def testIsBuiltin(self):
+    self.assertEqual(
+        tf_inspect.isbuiltin(TestDecoratedClass),
+        inspect.isbuiltin(TestDecoratedClass))
+    self.assertEqual(
+        tf_inspect.isbuiltin(test_decorated_function),
+        inspect.isbuiltin(test_decorated_function))
+    self.assertEqual(
+        tf_inspect.isbuiltin(test_undecorated_function),
+        inspect.isbuiltin(test_undecorated_function))
+    self.assertEqual(tf_inspect.isbuiltin(range), inspect.isbuiltin(range))
+    self.assertEqual(tf_inspect.isbuiltin(max), inspect.isbuiltin(max))
+
+  def testIsClass(self):
+    self.assertTrue(tf_inspect.isclass(TestDecoratedClass))
+    self.assertFalse(tf_inspect.isclass(test_decorated_function))
+
+  def testIsFunction(self):
+    self.assertTrue(tf_inspect.isfunction(test_decorated_function))
+    self.assertFalse(tf_inspect.isfunction(TestDecoratedClass))
+
+  def testIsMethod(self):
+    self.assertTrue(tf_inspect.ismethod(TestDecoratedClass().two))
+    self.assertFalse(tf_inspect.ismethod(test_decorated_function))
+
+  def testIsModule(self):
+    self.assertTrue(
+        tf_inspect.ismodule(inspect.getmodule(inspect.currentframe())))
+    self.assertFalse(tf_inspect.ismodule(test_decorated_function))
+
+  def testIsRoutine(self):
+    self.assertTrue(tf_inspect.isroutine(len))
+    self.assertFalse(tf_inspect.isroutine(TestDecoratedClass))
+
+  def testStack(self):
+    expected_stack = inspect.stack()
+    actual_stack = tf_inspect.stack()
+    self.assertEqual(len(expected_stack), len(actual_stack))
+    self.assertEqual(expected_stack[0][0], actual_stack[0][0])  # Frame object
+    self.assertEqual(expected_stack[0][1], actual_stack[0][1])  # Filename
+    self.assertEqual(expected_stack[0][2],
+                     actual_stack[0][2] - 1)  # Line number
+    self.assertEqual(expected_stack[0][3], actual_stack[0][3])  # Function name
+    self.assertEqual(expected_stack[1:], actual_stack[1:])
+
+  def testIsAnyTargetMethod(self):
+    class MyModule:
+
+      def f(self, a):
+        pass
+
+      def __call__(self):
+        pass
+
+    module = MyModule()
+    self.assertTrue(tf_inspect.isanytargetmethod(module))
+    f = module.f
+    self.assertTrue(tf_inspect.isanytargetmethod(f))
+    f = functools.partial(f, 1)
+    self.assertTrue(tf_inspect.isanytargetmethod(f))
+    f = test_decorator('tf_decorator1')(f)
+    self.assertTrue(tf_inspect.isanytargetmethod(f))
+    f = test_decorator('tf_decorator2')(f)
+    self.assertTrue(tf_inspect.isanytargetmethod(f))
+
+    class MyModule2:
+      pass
+    module = MyModule2()
+    self.assertFalse(tf_inspect.isanytargetmethod(module))
+    def f2(x):
+      del x
+    self.assertFalse(tf_inspect.isanytargetmethod(f2))
+    f2 = functools.partial(f2, 1)
+    self.assertFalse(tf_inspect.isanytargetmethod(f2))
+    f2 = test_decorator('tf_decorator1')(f2)
+    self.assertFalse(tf_inspect.isanytargetmethod(f2))
+    f2 = test_decorator('tf_decorator2')(f2)
+    self.assertFalse(tf_inspect.isanytargetmethod(f2))
+    self.assertFalse(tf_inspect.isanytargetmethod(lambda: None))
+    self.assertFalse(tf_inspect.isanytargetmethod(None))
+    self.assertFalse(tf_inspect.isanytargetmethod(1))
+
+
+class TfInspectGetCallArgsTest(test.TestCase):
+
+  def testReturnsEmptyWhenUnboundFuncHasNoParameters(self):
+
+    def empty():
+      pass
+
+    self.assertEqual({}, tf_inspect.getcallargs(empty))
+
+  def testClashingParameterNames(self):
+
+    def func(positional, func=1, func_and_positional=2, kwargs=3):
+      return positional, func, func_and_positional, kwargs
+
+    kwargs = {}
+    self.assertEqual(
+        tf_inspect.getcallargs(func, 0, **kwargs), {
+            'positional': 0,
+            'func': 1,
+            'func_and_positional': 2,
+            'kwargs': 3
+        })
+    kwargs = dict(func=4, func_and_positional=5, kwargs=6)
+    self.assertEqual(
+        tf_inspect.getcallargs(func, 0, **kwargs), {
+            'positional': 0,
+            'func': 4,
+            'func_and_positional': 5,
+            'kwargs': 6
+        })
+
+  def testUnboundFuncWithOneParamPositional(self):
+
+    def func(a):
+      return a
+
+    self.assertEqual({'a': 5}, tf_inspect.getcallargs(func, 5))
+
+  def testUnboundFuncWithTwoParamsPositional(self):
+
+    def func(a, b):
+      return (a, b)
+
+    self.assertEqual({'a': 10, 'b': 20}, tf_inspect.getcallargs(func, 10, 20))
+
+  def testUnboundFuncWithOneParamKeyword(self):
+
+    def func(a):
+      return a
+
+    self.assertEqual({'a': 5}, tf_inspect.getcallargs(func, a=5))
+
+  def testUnboundFuncWithTwoParamsKeyword(self):
+
+    def func(a, b):
+      return (a, b)
+
+    self.assertEqual({'a': 6, 'b': 7}, tf_inspect.getcallargs(func, a=6, b=7))
+
+  def testUnboundFuncWithOneParamDefault(self):
+
+    def func(a=13):
+      return a
+
+    self.assertEqual({'a': 13}, tf_inspect.getcallargs(func))
+
+  def testUnboundFuncWithOneParamDefaultOnePositional(self):
+
+    def func(a=0):
+      return a
+
+    self.assertEqual({'a': 1}, tf_inspect.getcallargs(func, 1))
+
+  def testUnboundFuncWithTwoParamsDefaultOnePositional(self):
+
+    def func(a=1, b=2):
+      return (a, b)
+
+    self.assertEqual({'a': 5, 'b': 2}, tf_inspect.getcallargs(func, 5))
+
+  def testUnboundFuncWithTwoParamsDefaultTwoPositional(self):
+
+    def func(a=1, b=2):
+      return (a, b)
+
+    self.assertEqual({'a': 3, 'b': 4}, tf_inspect.getcallargs(func, 3, 4))
+
+  def testUnboundFuncWithOneParamDefaultOneKeyword(self):
+
+    def func(a=1):
+      return a
+
+    self.assertEqual({'a': 3}, tf_inspect.getcallargs(func, a=3))
+
+  def testUnboundFuncWithTwoParamsDefaultOneKeywordFirst(self):
+
+    def func(a=1, b=2):
+      return (a, b)
+
+    self.assertEqual({'a': 3, 'b': 2}, tf_inspect.getcallargs(func, a=3))
+
+  def testUnboundFuncWithTwoParamsDefaultOneKeywordSecond(self):
+
+    def func(a=1, b=2):
+      return (a, b)
+
+    self.assertEqual({'a': 1, 'b': 4}, tf_inspect.getcallargs(func, b=4))
+
+  def testUnboundFuncWithTwoParamsDefaultTwoKeywords(self):
+
+    def func(a=1, b=2):
+      return (a, b)
+
+    self.assertEqual({'a': 3, 'b': 4}, tf_inspect.getcallargs(func, a=3, b=4))
+
+  def testBoundFuncWithOneParam(self):
+
+    class Test(object):
+
+      def bound(self):
+        pass
+
+    t = Test()
+    self.assertEqual({'self': t}, tf_inspect.getcallargs(t.bound))
+
+  def testBoundFuncWithManyParamsAndDefaults(self):
+
+    class Test(object):
+
+      def bound(self, a, b=2, c='Hello'):
+        return (a, b, c)
+
+    t = Test()
+    self.assertEqual({
+        'self': t,
+        'a': 3,
+        'b': 2,
+        'c': 'Goodbye'
+    }, tf_inspect.getcallargs(t.bound, 3, c='Goodbye'))
+
+  def testClassMethod(self):
+
+    class Test(object):
+
+      @classmethod
+      def test(cls, a, b=3, c='hello'):
+        return (a, b, c)
+
+    self.assertEqual({
+        'cls': Test,
+        'a': 5,
+        'b': 3,
+        'c': 'goodbye'
+    }, tf_inspect.getcallargs(Test.test, 5, c='goodbye'))
+
+  def testUsesOutermostDecoratorsArgSpec(self):
+
+    def func():
+      pass
+
+    def wrapper(*args, **kwargs):
+      return func(*args, **kwargs)
+
+    decorated = tf_decorator.make_decorator(
+        func,
+        wrapper,
+        decorator_argspec=tf_inspect.FullArgSpec(
+            args=['a', 'b', 'c'],
+            varargs=None,
+            kwonlyargs={},
+            defaults=(3, 'hello'),
+            kwonlydefaults=None,
+            varkw=None,
+            annotations=None))
+
+    self.assertEqual({
+        'a': 4,
+        'b': 3,
+        'c': 'goodbye'
+    }, tf_inspect.getcallargs(decorated, 4, c='goodbye'))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/tf_should_use.py

@@ -0,0 +1,311 @@
+# 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.
+# ==============================================================================
+"""Decorator that provides a warning if the wrapped object is never used."""
+import copy
+import sys
+import textwrap
+import traceback
+import types
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.util import tf_decorator
+
+
+class _TFShouldUseHelper(object):
+  """Object stored in TFShouldUse-wrapped objects.
+
+  When it is deleted it will emit a warning or error if its `sate` method
+  has not been called by time of deletion, and Tensorflow is not executing
+  eagerly or inside a tf.function (which use autodeps and resolve the
+  main issues this wrapper warns about).
+  """
+
+  def __init__(self, type_, repr_, stack_frame, error_in_function,
+               warn_in_eager):
+    self._type = type_
+    self._repr = repr_
+    self._stack_frame = stack_frame
+    self._error_in_function = error_in_function
+    if context.executing_eagerly():
+      # If warn_in_eager, sated == False.  Otherwise true.
+      self._sated = not warn_in_eager
+    elif ops.inside_function():
+      if error_in_function:
+        self._sated = False
+        ops.add_exit_callback_to_default_func_graph(
+            lambda: self._check_sated(raise_error=True))
+      else:
+        self._sated = True
+    else:
+      # TF1 graph building mode
+      self._sated = False
+
+  def sate(self):
+    self._sated = True
+    self._type = None
+    self._repr = None
+    self._stack_frame = None
+    self._logging_module = None
+
+  def _check_sated(self, raise_error):
+    """Check if the object has been sated."""
+    if self._sated:
+      return
+    creation_stack = ''.join(
+        [line.rstrip()
+         for line in traceback.format_stack(self._stack_frame, limit=5)])
+    if raise_error:
+      try:
+        raise RuntimeError(
+            'Object was never used (type {}): {}.  If you want to mark it as '
+            'used call its "mark_used()" method.  It was originally created '
+            'here:\n{}'.format(self._type, self._repr, creation_stack))
+      finally:
+        self.sate()
+    else:
+      tf_logging.error(
+          '==================================\n'
+          'Object was never used (type {}):\n{}\nIf you want to mark it as '
+          'used call its "mark_used()" method.\nIt was originally created '
+          'here:\n{}\n'
+          '=================================='
+          .format(self._type, self._repr, creation_stack))
+
+  def __del__(self):
+    self._check_sated(raise_error=False)
+
+
+def _new__init__(self, wrapped_value, tf_should_use_helper):
+  # pylint: disable=protected-access
+  self._tf_should_use_helper = tf_should_use_helper
+  self._tf_should_use_wrapped_value = wrapped_value
+
+
+def _new__setattr__(self, key, value):
+  if key in ('_tf_should_use_helper', '_tf_should_use_wrapped_value'):
+    return object.__setattr__(self, key, value)
+  return setattr(
+      object.__getattribute__(self, '_tf_should_use_wrapped_value'),
+      key, value)
+
+
+def _new__getattribute__(self, key):
+  if key not in ('_tf_should_use_helper', '_tf_should_use_wrapped_value'):
+    object.__getattribute__(self, '_tf_should_use_helper').sate()
+  if key in (
+      '_tf_should_use_wrapped_value',
+      '_tf_should_use_helper',
+      'mark_used',
+      '__setattr__',
+  ):
+    return object.__getattribute__(self, key)
+  return getattr(
+      object.__getattribute__(self, '_tf_should_use_wrapped_value'), key)
+
+
+def _new_mark_used(self, *args, **kwargs):
+  object.__getattribute__(self, '_tf_should_use_helper').sate()
+  try:
+    mu = object.__getattribute__(
+        object.__getattribute__(self, '_tf_should_use_wrapped_value'),
+        'mark_used')
+    return mu(*args, **kwargs)
+  except AttributeError:
+    pass
+
+OVERLOADABLE_OPERATORS = {
+    '__add__',
+    '__radd__',
+    '__sub__',
+    '__rsub__',
+    '__mul__',
+    '__rmul__',
+    '__div__',
+    '__rdiv__',
+    '__truediv__',
+    '__rtruediv__',
+    '__floordiv__',
+    '__rfloordiv__',
+    '__mod__',
+    '__rmod__',
+    '__lt__',
+    '__le__',
+    '__gt__',
+    '__ge__',
+    '__ne__',
+    '__eq__',
+    '__and__',
+    '__rand__',
+    '__or__',
+    '__ror__',
+    '__xor__',
+    '__rxor__',
+    '__getitem__',
+    '__pow__',
+    '__rpow__',
+    '__invert__',
+    '__neg__',
+    '__abs__',
+    '__matmul__',
+    '__rmatmul__',
+}
+
+
+_WRAPPERS = {}
+
+
+class ShouldUseWrapper(object):
+  pass
+
+
+def _get_wrapper(x, tf_should_use_helper):
+  """Create a wrapper for object x, whose class subclasses type(x).
+
+  The wrapper will emit a warning if it is deleted without any of its
+  properties being accessed or methods being called.
+
+  Args:
+    x: The instance to wrap.
+    tf_should_use_helper: The object that tracks usage.
+
+  Returns:
+    An object wrapping `x`, of type `type(x)`.
+  """
+  type_x = type(x)
+  memoized = _WRAPPERS.get(type_x, None)
+  if memoized:
+    return memoized(x, tf_should_use_helper)
+
+  # Make a copy of `object`
+  tx = copy.deepcopy(ShouldUseWrapper)
+  # Prefer using __orig_bases__, which preserve generic type arguments.
+  bases = getattr(tx, '__orig_bases__', tx.__bases__)
+
+  def set_body(ns):
+    ns.update(tx.__dict__)
+    return ns
+
+  copy_tx = types.new_class(tx.__name__, bases, exec_body=set_body)
+  copy_tx.__init__ = _new__init__
+  copy_tx.__getattribute__ = _new__getattribute__
+  for op in OVERLOADABLE_OPERATORS:
+    if hasattr(type_x, op):
+      setattr(copy_tx, op, getattr(type_x, op))
+
+  copy_tx.mark_used = _new_mark_used
+  copy_tx.__setattr__ = _new__setattr__
+  _WRAPPERS[type_x] = copy_tx
+
+  return copy_tx(x, tf_should_use_helper)
+
+
+def _add_should_use_warning(x, error_in_function=False, warn_in_eager=False):
+  """Wraps object x so that if it is never used, a warning is logged.
+
+  Args:
+    x: Python object.
+    error_in_function: Python bool.  If `True`, a `RuntimeError` is raised
+      if the returned value is never used when created during `tf.function`
+      tracing.
+    warn_in_eager: Python bool. If `True` raise warning if in Eager mode as well
+      as graph mode.
+
+  Returns:
+    An instance of `TFShouldUseWarningWrapper` which subclasses `type(x)`
+    and is a very shallow wrapper for `x` which logs access into `x`.
+  """
+  if x is None or (isinstance(x, list) and not x):
+    return x
+
+  if context.executing_eagerly() and not warn_in_eager:
+    return x
+
+  if ops.inside_function() and not error_in_function:
+    # We don't currently log warnings in tf.function calls, so just skip it.
+    return x
+
+  # Extract the current frame for later use by traceback printing.
+  try:
+    raise ValueError()
+  except ValueError:
+    stack_frame = sys.exc_info()[2].tb_frame.f_back
+
+  tf_should_use_helper = _TFShouldUseHelper(
+      type_=type(x),
+      repr_=repr(x),
+      stack_frame=stack_frame,
+      error_in_function=error_in_function,
+      warn_in_eager=warn_in_eager)
+
+  return _get_wrapper(x, tf_should_use_helper)
+
+
+def should_use_result(fn=None, warn_in_eager=False, error_in_function=False):
+  """Function wrapper that ensures the function's output is used.
+
+  If the output is not used, a `logging.error` is logged.  If
+  `error_in_function` is set, then a `RuntimeError` will be raised at the
+  end of function tracing if the output is not used by that point.
+
+  An output is marked as used if any of its attributes are read, modified, or
+  updated.  Examples when the output is a `Tensor` include:
+
+  - Using it in any capacity (e.g. `y = t + 0`, `sess.run(t)`)
+  - Accessing a property (e.g. getting `t.name` or `t.op`).
+  - Calling `t.mark_used()`.
+
+  Note, certain behaviors cannot be tracked - for these the object may not
+  be marked as used.  Examples include:
+
+  - `t != 0`.  In this case, comparison is done on types / ids.
+  - `isinstance(t, tf.Tensor)`.  Similar to above.
+
+  Args:
+    fn: The function to wrap.
+    warn_in_eager: Whether to create warnings in Eager as well.
+    error_in_function: Whether to raise an error when creating a tf.function.
+
+  Returns:
+    The wrapped function.
+  """
+  def decorated(fn):
+    """Decorates the input function."""
+    def wrapped(*args, **kwargs):
+      return _add_should_use_warning(fn(*args, **kwargs),
+                                     warn_in_eager=warn_in_eager,
+                                     error_in_function=error_in_function)
+    fn_doc = fn.__doc__ or ''
+    split_doc = fn_doc.split('\n', 1)
+    if len(split_doc) == 1:
+      updated_doc = fn_doc
+    else:
+      brief, rest = split_doc
+      updated_doc = '\n'.join([brief, textwrap.dedent(rest)])
+
+    note = ('\n\nNote: The output of this function should be used. If it is '
+            'not, a warning will be logged or an error may be raised. '
+            'To mark the output as used, call its .mark_used() method.')
+    return tf_decorator.make_decorator(
+        target=fn,
+        decorator_func=wrapped,
+        decorator_name='should_use_result',
+        decorator_doc=updated_doc + note)
+
+  if fn is not None:
+    return decorated(fn)
+  else:
+    return decorated

tensorflow/python/util/tf_should_use_test.py

@@ -0,0 +1,139 @@
+# 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.
+# ==============================================================================
+"""Unit tests for tf_should_use."""
+
+# pylint: disable=unused-import
+import contextlib
+import gc
+import sys
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.util import tf_should_use
+
+
+@contextlib.contextmanager
+def reroute_error():
+  """Temporarily reroute errors written to tf_logging.error into `captured`."""
+  with test.mock.patch.object(tf_should_use.tf_logging, 'error') as error:
+    yield error
+
+
+class TfShouldUseTest(test.TestCase):
+
+  def testAddShouldUseWarningWhenNotUsed(self):
+    c = constant_op.constant(0, name='blah0')
+    def in_this_function():
+      h = tf_should_use._add_should_use_warning(c, warn_in_eager=True)
+      del h
+    with reroute_error() as error:
+      in_this_function()
+    msg = '\n'.join(error.call_args[0])
+    self.assertIn('Object was never used', msg)
+    if not context.executing_eagerly():
+      self.assertIn('blah0:0', msg)
+    self.assertIn('in_this_function', msg)
+    self.assertFalse(gc.garbage)
+
+  def testAddShouldUseExceptionInEagerAndFunction(self):
+    def in_this_function():
+      c = constant_op.constant(0, name='blah0')
+      h = tf_should_use._add_should_use_warning(
+          c, warn_in_eager=True, error_in_function=True)
+      del h
+    if context.executing_eagerly():
+      with reroute_error() as error:
+        in_this_function()
+      msg = '\n'.join(error.call_args[0])
+      self.assertIn('Object was never used', msg)
+      self.assertIn('in_this_function', msg)
+      self.assertFalse(gc.garbage)
+
+    tf_fn_in_this_function = def_function.function(in_this_function)
+    with self.assertRaisesRegex(RuntimeError,
+                                r'Object was never used.*blah0:0'):
+      tf_fn_in_this_function()
+    self.assertFalse(gc.garbage)
+
+  def _testAddShouldUseWarningWhenUsed(self, fn, name):
+    c = constant_op.constant(0, name=name)
+    with reroute_error() as error:
+      h = tf_should_use._add_should_use_warning(c, warn_in_eager=True)
+      fn(h)
+      del h
+    error.assert_not_called()
+
+  def testAddShouldUseWarningWhenUsedWithAdd(self):
+    def add(h):
+      _ = h + 1
+    self._testAddShouldUseWarningWhenUsed(add, name='blah_add')
+    gc.collect()
+    self.assertFalse(gc.garbage)
+
+  def testAddShouldUseWarningWhenUsedWithGetShape(self):
+    def get_shape(h):
+      _ = h.shape
+    self._testAddShouldUseWarningWhenUsed(get_shape, name='blah_get_name')
+    gc.collect()
+    self.assertFalse(gc.garbage)
+
+  def testShouldUseResult(self):
+    @tf_should_use.should_use_result(warn_in_eager=True)
+    def return_const(value):
+      return constant_op.constant(value, name='blah2')
+    with reroute_error() as error:
+      return_const(0.0)
+    msg = '\n'.join(error.call_args[0])
+    self.assertIn('Object was never used', msg)
+    if not context.executing_eagerly():
+      self.assertIn('blah2:0', msg)
+    self.assertIn('return_const', msg)
+    gc.collect()
+    self.assertFalse(gc.garbage)
+
+  def testShouldUseResultWhenNotReallyUsed(self):
+    @tf_should_use.should_use_result(warn_in_eager=True)
+    def return_const(value):
+      return constant_op.constant(value, name='blah3')
+    with reroute_error() as error:
+      with self.cached_session():
+        return_const(0.0)
+        # Creating another op and executing it does not mark the
+        # unused op as being "used".
+        v = constant_op.constant(1.0, name='meh')
+        self.evaluate(v)
+    msg = '\n'.join(error.call_args[0])
+    self.assertIn('Object was never used', msg)
+    if not context.executing_eagerly():
+      self.assertIn('blah3:0', msg)
+    self.assertIn('return_const', msg)
+    gc.collect()
+    self.assertFalse(gc.garbage)
+
+  # Tests that mark_used is available in the API.
+  def testMarkUsed(self):
+    @tf_should_use.should_use_result(warn_in_eager=True)
+    def return_const(value):
+      return constant_op.constant(value, name='blah3')
+
+    with self.cached_session():
+      return_const(0.0).mark_used()
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/util/tf_stack.cc

@@ -0,0 +1,371 @@
+/* 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.
+*/
+
+// We extract stack traces in Python using the logic in tf_stack.cc, which
+// stores a list of PyCodeObject*. Such stack trace extraction is really fast.
+//
+// We store the retrieved stack trace within the Node object directly. Then
+// whenever the graph is instantiated/copies, we copy the stack trace with it.
+// Since the graph instantiation goes through the protobuf roundtrip, we store
+// the original stack traces mapping attached in FunctionLibraryDefinition.
+
+// clang-format off
+// These headers must be at the top, before including Python.h header
+// Otherwise, we get C2039 on MSVC due to 'copysign'
+#include "absl/strings/str_cat.h"
+#include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil
+#include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil
+#include "pybind11/complex.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "pybind11/stl_bind.h"  // from @pybind11
+// clang-format on
+#include <frameobject.h>
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/hash/hash.h"
+#include "absl/strings/str_format.h"
+#include "absl/types/span.h"
+#include "tensorflow/core/graph/graph_debug_info_builder.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/stack_frame.h"
+#include "tensorflow/core/util/managed_stack_trace.h"
+#include "tensorflow/python/util/stack_trace.h"
+#include "tsl/platform/mutex.h"
+
+struct StackFrame;  // Forward declaration.
+struct StackTrace;
+
+PYBIND11_MAKE_OPAQUE(std::vector<StackFrame>);
+PYBIND11_MAKE_OPAQUE(StackTrace);
+
+namespace tensorflow {
+
+namespace {
+
+namespace py = pybind11;
+
+using StringSet = absl::flat_hash_set<std::string>;
+
+// Python wrapper for a SourceMap.
+class PyBindSourceMap {
+ public:
+  PyBindSourceMap() : source_map_(std::make_shared<SourceMap>()) {}
+
+  // Shares ownership with whoever captures traces in the scope of this map.
+  std::shared_ptr<SourceMap> source_map_;
+};
+
+// Python wrapper for a FileSet.
+class PyBindFileSet {
+ public:
+  PyBindFileSet() : file_set_(std::make_shared<StringSet>()) {}
+
+  // Shares ownership with whoever captures traces in the scope of this set.
+  std::shared_ptr<StringSet> file_set_;
+};
+
+// Simple caching wrapper around a captured stack trace.
+//
+// When required, stacks are computed and cached as a `FrozenStackTrace`.
+class StackTraceWrapper : public AbstractStackTrace {
+ public:
+  StackTraceWrapper(const std::shared_ptr<StackTrace>& captured,
+                    const std::shared_ptr<SourceMap>& source_map,
+                    const std::shared_ptr<StringSet>& filter, int stacklevel)
+      : captured_(captured),
+        source_map_(source_map),
+        filter_(filter),
+        stacklevel_(stacklevel) {}
+
+  ~StackTraceWrapper() override {
+    PyGILState_STATE state = PyGILState_Ensure();
+    captured_.reset();
+    source_map_.reset();
+    filter_.reset();
+    PyGILState_Release(state);
+  }
+
+  StackTraceWrapper(StackTraceWrapper&& rhs) = default;
+  StackTraceWrapper& operator=(StackTraceWrapper&& rhs) = default;
+
+  static std::unique_ptr<StackTraceWrapper> ExtractStack(
+      const std::shared_ptr<SourceMap>& source_map,
+      const std::shared_ptr<StringSet>& filter, int stacklevel) {
+    return std::make_unique<StackTraceWrapper>(StackTrace::Capture(-1),
+                                               source_map, filter, stacklevel);
+  }
+
+  absl::Span<const StackFrame> ToFrames() const override {
+    ComputeFrozen();
+    return cache_->ToFrames();
+  }
+
+  std::vector<StackFrame> ToUncachedFrames() const override {
+    std::vector<StackFrame> frames = captured_->ToStackFrames(
+        *source_map_, [&](const char* f) { return StackTraceFiltering(f); },
+        /*reverse_traversal=*/false, /*limit=*/-1);
+
+    // Drop last stack frames.
+    int newsize = frames.size() - stacklevel_;
+    if (newsize < 0) {
+      newsize = 0;
+    }
+    frames.resize(newsize);
+
+    return frames;
+  }
+
+  std::vector<StackFrame> GetUserFrames(int limit) const override {
+    ComputeFrozen();
+    return cache_->GetUserFrames(limit);
+  }
+
+  StackFrame LastUserFrame() const override {
+    ComputeFrozen();
+    return cache_->LastUserFrame();
+  }
+
+  std::string ToString(const TracePrintingOptions& opts) const override {
+    ComputeFrozen();
+    return cache_->ToString(opts);
+  }
+
+ private:
+  void ComputeFrozen() const {
+    tsl::mutex_lock lock(mu_);
+    if (cache_ != nullptr) {
+      return;
+    }
+
+    std::vector<StackFrame> frames = ToUncachedFrames();
+
+    std::vector<StackFrame> user_frames = captured_->ToStackFrames(
+        *source_map_,
+        [&](const char* file_name) {
+          return StackTraceFiltering(file_name) ||
+                 IsInternalFrameForFilename(file_name);
+        },
+        /*reverse_traversal=*/true,
+        /*limit=*/-1);
+    // ensure we use the original (outermost first) ordering.
+    absl::c_reverse(user_frames);
+
+    cache_ = std::make_unique<FrozenStackTrace>(frames, user_frames);
+  }
+
+  bool StackTraceFiltering(const char* file_name) const {
+    return filter_->contains(file_name);
+  }
+
+  mutable mutex mu_;
+  mutable std::unique_ptr<FrozenStackTrace> cache_;
+  std::shared_ptr<const StackTrace> captured_;
+  std::shared_ptr<SourceMap> source_map_;
+  std::shared_ptr<StringSet> filter_;
+  int stacklevel_;
+};
+
+}  // namespace
+
+PYBIND11_MODULE(_tf_stack, m, pybind11::mod_gil_not_used()) {
+  pybind11::google::ImportStatusModule();
+
+  py::class_<PyBindSourceMap>(m, "PyBindSourceMap")
+      .def(py::init())
+      .def("update_to",
+           [](const PyBindSourceMap& self, const py::tuple& source_map) {
+             self.source_map_->clear();
+             for (const auto& item : source_map) {
+               const auto& tuple_item = py::cast<py::tuple>(item);
+
+               const auto& key = py::cast<py::tuple>(tuple_item[0]);
+               std::string&& k_filename = py::cast<std::string>(key[0]);
+               int k_lineno = py::cast<int>(key[1]);
+
+               const auto& value = py::cast<py::tuple>(tuple_item[1]);
+               std::string&& v_filename = py::cast<std::string>(value[0]);
+               int v_lineno = py::cast<int>(value[1]);
+               const auto& function_name_val = value[2];
+               std::string&& v_function_name =
+                   function_name_val.is_none()
+                       ? ""
+                       : py::cast<std::string>(function_name_val);
+
+               self.source_map_->emplace(
+                   SourceLoc{k_filename, k_lineno},
+                   StackFrame({v_filename, v_lineno, v_function_name}));
+             }
+           });
+
+  py::class_<PyBindFileSet>(m, "PyBindFileSet")
+      .def(py::init())
+      .def("update_to", [](const PyBindFileSet& self, const py::set& file_set) {
+        self.file_set_->clear();
+        for (const auto& item : file_set) {
+          self.file_set_->insert(py::cast<std::string>(item));
+        }
+      });
+
+  py::class_<GraphDebugInfoBuilder>(m, "GraphDebugInfoBuilder")
+      .def(py::init())
+      .def(
+          "AppendGraphDebugInfo",
+          [](GraphDebugInfoBuilder& self, std::string fn_name,
+             py::bytes debug_info_str) {
+            return self.AppendGraphDebugInfoStr(fn_name, debug_info_str);
+          },
+          py::arg("prefix"), py::arg("debug_info"))
+      .def(
+          "AccumulateStackTrace",
+          [](GraphDebugInfoBuilder& self, std::string function, std::string op,
+             const AbstractStackTrace& trace) {
+            std::string key = absl::StrCat(op, "@", function);
+            self.AccumulateStackTrace(
+                std::make_shared<FrozenStackTrace>(trace.ToFrames()), key);
+          },
+          py::arg("function"), py::arg("op"), py::arg("trace"))
+      .def("Build", [](GraphDebugInfoBuilder& self) -> py::bytes {
+        return py::bytes(self.ToGraphDebugInfoStr());
+      });
+
+  py::class_<StackFrame>(m, "StackFrame")
+      .def_property_readonly(
+          "filename",
+          [](const StackFrame& self) { return py::str(self.file_name); })
+      .def_property_readonly(
+          "lineno",
+          [](const StackFrame& self) { return py::int_(self.line_number); })
+      .def_property_readonly(
+          "name",
+          [](const StackFrame& self) { return py::str(self.function_name); })
+      .def_property_readonly("line", [](const StackFrame& self) { return ""; })
+      .def("__eq__", &StackFrame::operator==)
+      .def("__ne__", &StackFrame::operator!=)
+      .def("__hash__",
+           [](const StackFrame& self) {
+             return absl::Hash<std::tuple<std::string, int, std::string>>()(
+                 std::make_tuple(self.file_name, self.line_number,
+                                 self.function_name));
+           })
+      .def("__getitem__",
+           [](const StackFrame& self, const py::object& index) -> py::object {
+             return py::make_tuple(
+                 py::str(self.file_name), py::int_(self.line_number),
+                 py::str(self.function_name), py::str(""))[index];
+           })
+      .def("__iter__",
+           [](const StackFrame& self) -> py::iterator {
+             return py::iter(py::make_tuple(
+                 py::str(self.file_name), py::int_(self.line_number),
+                 py::str(self.function_name), py::str("")));
+           })
+      .def("__repr__",
+           [](const StackFrame& self) -> py::str {
+             return absl::StrFormat("File \"%s\", line %d, in %s",
+                                    self.file_name, self.line_number,
+                                    py::str(self.function_name));
+           })
+      .def("__len__", [](const StackFrame&) { return 4; });
+
+  py::class_<AbstractStackTrace, std::shared_ptr<AbstractStackTrace>>(
+      m, "StackTrace")
+      .def(
+          "__getitem__",
+          [](const AbstractStackTrace& self, py::ssize_t index) -> StackFrame {
+            absl::Span<const StackFrame> frames = self.ToFrames();
+            const size_t eff_index =
+                index < 0 ? frames.size() + index : static_cast<size_t>(index);
+            if (eff_index >= frames.size()) {
+              throw py::index_error();
+            }
+            return frames[eff_index];
+          },
+          py::return_value_policy::take_ownership)
+      .def(
+          "__getitem__",
+          [](const AbstractStackTrace& self,
+             py::slice slice) -> std::shared_ptr<AbstractStackTrace> {
+            absl::Span<const StackFrame> frames = self.ToFrames();
+            py::ssize_t start, stop, step, slicelength;
+            if (!slice.compute(frames.size(), &start, &stop, &step,
+                               &slicelength)) {
+              throw py::error_already_set();
+            }
+            if (step == 1) {
+              return std::make_shared<FrozenStackTrace>(
+                  frames.subspan(start, slicelength));
+            }
+            std::vector<StackFrame> out;
+            if (slicelength > 0) {
+              out.reserve(slicelength);
+            }
+            // Python slices allow negative indexing.
+            for (py::ssize_t i = start, count = 0; count < slicelength;
+                 i += step, ++count) {
+              out.push_back(frames[static_cast<size_t>(i)]);
+            }
+            return std::make_shared<FrozenStackTrace>(out);
+          },
+          py::return_value_policy::take_ownership)
+      .def(
+          "__len__",
+          [](const AbstractStackTrace& self) { return self.ToFrames().size(); })
+      .def("__eq__",
+           [](const AbstractStackTrace& self, const AbstractStackTrace& other) {
+             return self.ToFrames() == other.ToFrames();
+           })
+      .def("__hash__",
+           [](const AbstractStackTrace& self) {
+             return py::hash(py::str(self.ToString({})));
+           })
+      .def(
+          "get_user_frames",
+          [](const AbstractStackTrace& self)
+              -> std::shared_ptr<AbstractStackTrace> {
+            return std::make_shared<FrozenStackTrace>(self.GetUserFrames(-1));
+          },
+          "Returns the non-framework frames as a new trace object.")
+      .def(
+          "last_user_frame",
+          [](const AbstractStackTrace& self) { return self.LastUserFrame(); },
+          "Returns the last non-framework frame.")
+      .def("__repr__",
+           [](const AbstractStackTrace& self) { return self.ToString({}); });
+
+  m.def(
+      "extract_stack",
+      [](const PyBindSourceMap& source_map, const PyBindFileSet& file_set,
+         int stacklevel) -> std::shared_ptr<AbstractStackTrace> {
+        return StackTraceWrapper::ExtractStack(source_map.source_map_,
+                                               file_set.file_set_, stacklevel);
+      },
+      py::arg("source_map"), py::arg("file_set"), py::arg("stacklevel") = 1,
+      py::return_value_policy::take_ownership);
+
+  m.def(
+      "LoadTracesFromDebugInfo",
+      [](py::bytes data) { return LoadTracesFromDebugInfoStr(data); },
+      py::arg("debug_info_proto"));
+}
+
+}  // namespace tensorflow

tensorflow/python/util/tf_stack.py

@@ -0,0 +1,187 @@
+# 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.
+# ==============================================================================
+"""Functions used to extract and analyze stacks.  Faster than Python libs."""
+# pylint: disable=g-bad-name
+import collections
+import inspect
+import threading
+
+from tensorflow.core.framework import graph_debug_info_pb2
+from tensorflow.python.util import _tf_stack
+
+# Generally such lookups should be done using `threading.local()`. See
+# https://blogs.gnome.org/jamesh/2008/06/11/tls-python/ for a detailed
+# explanation of why. However the transform stacks are expected to be empty
+# when a thread is joined, so reusing the key does not introduce a correctness
+# issue. Moreover, get_ident is faster than storing and retrieving a unique
+# key in a thread local store.
+_get_thread_key = threading.get_ident
+
+
+# TODO(mdan): Move these to C++ as well.
+# Moving to C++ can further avoid extra copies made by get_effective_map.
+_source_mapper_stacks = collections.defaultdict(lambda: [SentinelMapper()])
+_source_filter_stacks = collections.defaultdict(lambda: [SentinelFilter()])
+
+
+class StackTraceTransform(object):
+  """Base class for stack trace transformation functions."""
+
+  _stack_dict = None  # Subclasses should override
+  _thread_key = None
+
+  def __enter__(self):
+    # Any given instance is assumed to be used by a single thread, which reduces
+    # expensive thread local lookups.
+    if self._thread_key is None:
+      self._thread_key = _get_thread_key()
+    else:
+      assert self._thread_key == _get_thread_key(), 'Shared across threads?'
+
+    stack = self._stack_dict[self._thread_key]
+    self.parent = stack[-1]
+    stack.append(self)
+    self.update()
+    return self
+
+  def __exit__(self, unused_type, unused_value, unused_traceback):
+    top = self._stack_dict[self._thread_key].pop()
+    assert top is self, 'Concurrent access?'
+
+  def update(self):
+    raise NotImplementedError('subclasses need to override this')
+
+
+class StackTraceMapper(StackTraceTransform):
+  """Allows remapping traceback information to different source code."""
+  _stack_dict = _source_mapper_stacks
+
+  def __init__(self):
+    self.internal_map = _tf_stack.PyBindSourceMap()
+
+  def update(self):
+    self.internal_map.update_to(tuple(self.get_effective_source_map().items()))
+
+  def get_effective_source_map(self):
+    """Returns a map (filename, lineno) -> (filename, lineno, function_name)."""
+    raise NotImplementedError('subclasses need to override this')
+
+
+EMPTY_DICT = {}
+
+
+class SentinelMapper(StackTraceMapper):
+
+  def get_effective_source_map(self):
+    return EMPTY_DICT
+
+
+class StackTraceFilter(StackTraceTransform):
+  """Allows filtering traceback information by removing superfluous frames."""
+  _stack_dict = _source_filter_stacks
+
+  def __init__(self):
+    self.internal_set = _tf_stack.PyBindFileSet()
+
+  def update(self):
+    self.internal_set.update_to(set(self.get_filtered_filenames()))
+
+  def get_filtered_filenames(self):
+    raise NotImplementedError('subclasses need to override this')
+
+
+EMPTY_SET = frozenset()
+
+
+class SentinelFilter(StackTraceFilter):
+
+  def get_filtered_filenames(self):
+    return EMPTY_SET
+
+
+class CurrentModuleFilter(StackTraceFilter):
+  """Filters stack frames from the module where this is used (best effort)."""
+
+  def __init__(self):
+    super().__init__()
+    filter_filename = None
+    outer_f = None
+    f = inspect.currentframe()
+    try:
+      if f is not None:
+        # The current frame is __init__. The first outer frame should be the
+        # caller.
+        outer_f = f.f_back
+        if outer_f is not None:
+          filter_filename = inspect.getfile(outer_f)
+      self._filename = filter_filename
+      # This may be called repeatedly: once on entry by the superclass, then by
+      # each child context manager.
+      self._cached_set = None
+    finally:
+      # Avoid reference cycles, see:
+      # https://docs.python.org/3.7/library/inspect.html#the-interpreter-stack
+      del f
+      del outer_f
+
+  def get_filtered_filenames(self):
+    if self._cached_set is not None:
+      return self._cached_set
+
+    filtered_filenames = frozenset((self._filename,))
+    if self.parent is not None:
+      filtered_filenames |= self.parent.get_filtered_filenames()
+    self._cached_set = filtered_filenames
+    return filtered_filenames
+
+
+def extract_stack(stacklevel=1):
+  """An eager-friendly alternative to traceback.extract_stack.
+
+  Args:
+    stacklevel: number of initial frames to skip when producing the stack.
+
+  Returns:
+    A list-like FrameSummary containing StackFrame-like objects, which are
+    namedtuple-like objects with the following fields: filename, lineno, name,
+    line, meant to masquerade as traceback.FrameSummary objects.
+  """
+  thread_key = _get_thread_key()
+  return _tf_stack.extract_stack(
+      _source_mapper_stacks[thread_key][-1].internal_map,
+      _source_filter_stacks[thread_key][-1].internal_set,
+      stacklevel,
+  )
+
+
+def LoadTracesFromDebugInfo(debug_info):
+  return _tf_stack.LoadTracesFromDebugInfo(debug_info.SerializeToString())
+
+
+class GraphDebugInfoBuilder(_tf_stack.GraphDebugInfoBuilder):
+
+  def AppendGraphDebugInfo(self, fn_name, fn_debug_info):
+    debug_info_str = fn_debug_info.SerializeToString()
+    super().AppendGraphDebugInfo(fn_name, debug_info_str)
+
+  def Build(self):
+    debug_info_str = super().Build()
+    debug_info = graph_debug_info_pb2.GraphDebugInfo()
+    debug_info.ParseFromString(debug_info_str)
+    return debug_info
+
+
+StackSummary = _tf_stack.StackTrace
+FrameSummary = _tf_stack.StackFrame

tensorflow/python/util/tf_stack_test.py

@@ -0,0 +1,149 @@
+# 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 functions used to extract and analyze stacks."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.util import tf_stack
+
+
+class TFStackTest(test.TestCase):
+
+  def testFrameSummaryEquality(self):
+    frames1 = tf_stack.extract_stack()
+    frames2 = tf_stack.extract_stack()
+
+    self.assertNotEqual(frames1[0], frames1[1])
+    self.assertEqual(frames1[0], frames1[0])
+    self.assertEqual(frames1[0], frames2[0])
+
+  def testFrameSummaryEqualityAndHash(self):
+    # Both defined on the same line to produce identical stacks.
+    frame1, frame2 = tf_stack.extract_stack(), tf_stack.extract_stack()
+    self.assertEqual(len(frame1), len(frame2))
+    for f1, f2 in zip(frame1, frame2):
+      self.assertEqual(f1, f2)
+      self.assertEqual(hash(f1), hash(f1))
+      self.assertEqual(hash(f1), hash(f2))
+    self.assertEqual(frame1, frame2)
+    self.assertEqual(hash(tuple(frame1)), hash(tuple(frame2)))
+
+  def testLastUserFrame(self):
+    trace = tf_stack.extract_stack()
+    frame = trace.last_user_frame()
+    self.assertRegex(repr(frame), 'testLastUserFrame')
+
+  def testGetUserFrames(self):
+
+    def func():
+      trace = tf_stack.extract_stack()  # COMMENT
+      frames = list(trace.get_user_frames())
+      return frames
+
+    frames = func()  # CALLSITE
+
+    self.assertRegex(repr(frames[-1]), 'func')
+    self.assertRegex(repr(frames[-2]), 'testGetUserFrames')
+
+  def testGetItem(self):
+    def func(n):
+      if n == 0:
+        return tf_stack.extract_stack()  # COMMENT
+      else:
+        return func(n - 1)
+
+    trace = func(5)
+    self.assertIn('func', repr(trace[-1]))
+
+    with self.assertRaises(IndexError):
+      _ = trace[-len(trace) - 1]
+
+    with self.assertRaises(IndexError):
+      _ = trace[len(trace)]
+
+  def testSlice(self):
+    def func(n):
+      if n == 0:
+        return tf_stack.extract_stack()
+      else:
+        return func(n - 1)
+
+    trace = func(5)
+
+    # Test full slice (step=1)
+    sliced_full = trace[:]
+    self.assertEqual(len(sliced_full), len(trace))
+
+    # Test partial slice (step=1)
+    if len(trace) > 2:
+      sliced_partial = trace[1:3]
+      self.assertEqual(len(sliced_partial), 2)
+
+    # Test slice with step != 1
+    sliced_step = trace[::2]
+    self.assertEqual(len(sliced_step), (len(trace) + 1) // 2)
+
+    # Test negative step
+    sliced_reverse = trace[::-1]
+    self.assertEqual(len(sliced_reverse), len(trace))
+
+    # Test empty slice
+    sliced_empty = trace[5:1]
+    self.assertEqual(len(sliced_empty), 0)
+
+  def testSourceMap(self):
+    source_map = tf_stack._tf_stack.PyBindSourceMap()
+
+    def func(n):
+      if n == 0:
+        return tf_stack._tf_stack.extract_stack(
+            source_map, tf_stack._tf_stack.PyBindFileSet()
+        )
+      else:
+        return func(n - 1)
+
+    trace = func(5)
+    source_map.update_to((
+        (
+            (trace[0].filename, trace[0].lineno),
+            ("filename", 42, "function_name"),
+        ),
+    ))
+    trace = list(func(5))
+    self.assertEqual(
+        str(trace[0]), 'File "filename", line 42, in function_name'
+    )
+
+  def testStackTraceBuilder(self):
+    stack1 = tf_stack.extract_stack()
+    stack2 = tf_stack.extract_stack()
+    stack3 = tf_stack.extract_stack()
+
+    builder = tf_stack.GraphDebugInfoBuilder()
+    builder.AccumulateStackTrace('func1', 'node1', stack1)
+    builder.AccumulateStackTrace('func2', 'node2', stack2)
+    builder.AccumulateStackTrace('func3', 'node3', stack3)
+    debug_info = builder.Build()
+
+    trace_map = tf_stack.LoadTracesFromDebugInfo(debug_info)
+    self.assertSameElements(
+        trace_map.keys(), ['node1@func1', 'node2@func2', 'node3@func3']
+    )
+
+    for trace in trace_map.values():
+      self.assertRegex(repr(trace), 'tf_stack_test.py', trace)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/util/tfprof_wrapper.cc

@@ -0,0 +1,45 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/profiler/internal/print_model_analysis.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_tfprof, m) {
+  m.def("PrintModelAnalysis",
+        [](const std::string* graph, const std::string* run_meta,
+           const std::string* op_log, const std::string* command,
+           const std::string* options) {
+          std::string temp = tensorflow::tfprof::PrintModelAnalysis(
+              graph, run_meta, op_log, command, options);
+          return py::bytes(temp);
+        });
+  m.def("NewProfiler", &tensorflow::tfprof::NewProfiler);
+  m.def("ProfilerFromFile", &tensorflow::tfprof::ProfilerFromFile);
+  m.def("DeleteProfiler", &tensorflow::tfprof::DeleteProfiler);
+  m.def("AddStep", &tensorflow::tfprof::AddStep);
+  m.def("SerializeToString", []() {
+    std::string temp = tensorflow::tfprof::SerializeToString();
+    return py::bytes(temp);
+  });
+  m.def("WriteProfile", &tensorflow::tfprof::WriteProfile);
+  m.def("Profile", [](const std::string* command, const std::string* options) {
+    std::string temp = tensorflow::tfprof::Profile(command, options);
+    return py::bytes(temp);
+  });
+}

tensorflow/python/util/traceback_utils.py

@@ -0,0 +1,171 @@
+# 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.
+# ==============================================================================
+"""Utilities related to TensorFlow exception stack trace prettifying."""
+
+import os
+import sys
+import threading
+import traceback
+import types
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util.tf_export import tf_export
+
+
+_ENABLE_TRACEBACK_FILTERING = threading.local()
+_EXCLUDED_PATHS = (
+    os.path.abspath(os.path.join(__file__, '..', '..')),
+)
+
+
+@tf_export('debugging.is_traceback_filtering_enabled')
+def is_traceback_filtering_enabled():
+  """Check whether traceback filtering is currently enabled.
+
+  See also `tf.debugging.enable_traceback_filtering()` and
+  `tf.debugging.disable_traceback_filtering()`. Note that filtering out
+  internal frames from the tracebacks of exceptions raised by TensorFlow code
+  is the default behavior.
+
+  Returns:
+    True if traceback filtering is enabled
+    (e.g. if `tf.debugging.enable_traceback_filtering()` was called),
+    and False otherwise (e.g. if `tf.debugging.disable_traceback_filtering()`
+    was called).
+  """
+  value = getattr(_ENABLE_TRACEBACK_FILTERING, 'value', True)
+  return value
+
+
+@tf_export('debugging.enable_traceback_filtering')
+def enable_traceback_filtering():
+  """Enable filtering out TensorFlow-internal frames in exception stack traces.
+
+  Raw TensorFlow stack traces involve many internal frames, which can be
+  challenging to read through, while not being actionable for end users.
+  By default, TensorFlow filters internal frames in most exceptions that it
+  raises, to keep stack traces short, readable, and focused on what's
+  actionable for end users (their own code).
+
+  If you have previously disabled traceback filtering via
+  `tf.debugging.disable_traceback_filtering()`, you can re-enable it via
+  `tf.debugging.enable_traceback_filtering()`.
+
+  Raises:
+    RuntimeError: If Python version is not at least 3.7.
+  """
+  if sys.version_info.major != 3 or sys.version_info.minor < 7:
+    raise RuntimeError(
+        f'Traceback filtering is only available with Python 3.7 or higher. '
+        f'This Python version: {sys.version}')
+  global _ENABLE_TRACEBACK_FILTERING
+  _ENABLE_TRACEBACK_FILTERING.value = True
+
+
+@tf_export('debugging.disable_traceback_filtering')
+def disable_traceback_filtering():
+  """Disable filtering out TensorFlow-internal frames in exception stack traces.
+
+  Raw TensorFlow stack traces involve many internal frames, which can be
+  challenging to read through, while not being actionable for end users.
+  By default, TensorFlow filters internal frames in most exceptions that it
+  raises, to keep stack traces short, readable, and focused on what's
+  actionable for end users (their own code).
+
+  Calling `tf.debugging.disable_traceback_filtering` disables this filtering
+  mechanism, meaning that TensorFlow exceptions stack traces will include
+  all frames, in particular TensorFlow-internal ones.
+
+  **If you are debugging a TensorFlow-internal issue, you need to call
+  `tf.debugging.disable_traceback_filtering`**.
+  To re-enable traceback filtering afterwards, you can call
+  `tf.debugging.enable_traceback_filtering()`.
+  """
+  global _ENABLE_TRACEBACK_FILTERING
+  _ENABLE_TRACEBACK_FILTERING.value = False
+
+
+def include_frame(fname):
+  """Determines whether a frame should be included in filtered tracebacks.
+
+  This function checks if the given filename contains any of the excluded
+  paths. Frames from excluded paths (typically TensorFlow internal code) are
+  filtered out from stack traces to make them more readable for end users.
+
+  Args:
+    fname: A string representing the filename of the frame to check.
+
+  Returns:
+    True if the frame should be included in the traceback (i.e., it's not
+    from an excluded path), False otherwise.
+  """
+
+  for exclusion in _EXCLUDED_PATHS:
+    if exclusion in fname:
+      return False
+  return True
+
+
+def _process_traceback_frames(tb):
+  new_tb = None
+  tb_list = list(traceback.walk_tb(tb))
+  for f, line_no in reversed(tb_list):
+    if include_frame(f.f_code.co_filename):
+      new_tb = types.TracebackType(new_tb, f, f.f_lasti, line_no)
+  if new_tb is None and tb_list:
+    f, line_no = tb_list[-1]
+    new_tb = types.TracebackType(new_tb, f, f.f_lasti, line_no)
+  return new_tb
+
+
+def filter_traceback(fn):
+  """Decorator to filter out TF-internal stack trace frames in exceptions.
+
+  Raw TensorFlow stack traces involve many internal frames, which can be
+  challenging to read through, while not being actionable for end users.
+  By default, TensorFlow filters internal frames in most exceptions that it
+  raises, to keep stack traces short, readable, and focused on what's
+  actionable for end users (their own code).
+
+  Arguments:
+    fn: The function or method to decorate. Any exception raised within the
+      function will be reraised with its internal stack trace frames filtered
+      out.
+
+  Returns:
+    Decorated function or method.
+  """
+  if sys.version_info.major != 3 or sys.version_info.minor < 7:
+    return fn
+
+  def error_handler(*args, **kwargs):
+    try:
+      if not is_traceback_filtering_enabled():
+        return fn(*args, **kwargs)
+    except NameError:
+      # In some very rare cases,
+      # `is_traceback_filtering_enabled` (from the outer scope) may not be
+      # accessible from inside this function
+      return fn(*args, **kwargs)
+
+    filtered_tb = None
+    try:
+      return fn(*args, **kwargs)
+    except Exception as e:
+      filtered_tb = _process_traceback_frames(e.__traceback__)
+      raise e.with_traceback(filtered_tb) from None
+    finally:
+      del filtered_tb
+
+  return tf_decorator.make_decorator(fn, error_handler)

tensorflow/python/util/traceback_utils_test.py

@@ -0,0 +1,119 @@
+# 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 traceback_utils."""
+
+import sys
+import traceback
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.util import traceback_utils
+
+
+class TracebackUtilsTest(test.TestCase):
+
+  def assert_trace_line_count(self, fn, count, filtering_enabled=True):
+    trace_line_count = -1
+    if filtering_enabled:
+      traceback_utils.enable_traceback_filtering()
+    else:
+      traceback_utils.disable_traceback_filtering()
+    self.assertEqual(
+        traceback_utils.is_traceback_filtering_enabled(), filtering_enabled)
+    try:
+      fn()
+    except Exception as e:  # pylint: disable=broad-except
+      # We must count lines rather than frames because autograph transforms
+      # stack frames into a single large string
+      trace = '\n'.join(traceback.format_tb(e.__traceback__))
+      trace_line_count = len(trace.split('\n'))
+
+    self.assertGreater(trace_line_count, 0)
+
+    if filtering_enabled:
+      if sys.version_info >= (3, 13):
+        self.assertLessEqual(trace_line_count, count)
+      else:
+        self.assertLess(trace_line_count, count)
+    else:
+      self.assertGreater(trace_line_count, count)
+
+  def test_eager_add(self):
+
+    def fn():
+      x = array_ops.zeros((2, 3))
+      y = array_ops.zeros((2, 4))
+      _ = x + y
+
+    self.assert_trace_line_count(fn, count=15, filtering_enabled=True)
+    self.assert_trace_line_count(fn, count=25, filtering_enabled=False)
+
+  def test_tfn_add(self):
+    @def_function.function
+    def fn():
+      x = array_ops.zeros((2, 3))
+      y = array_ops.zeros((2, 4))
+      return x + y
+
+    self.assert_trace_line_count(fn, count=10, filtering_enabled=True)
+    self.assert_trace_line_count(fn, count=25, filtering_enabled=False)
+
+  def test_tfn_div(self):
+    @def_function.function
+    def wrapped_fn(x):
+      return x / 0.
+
+    def fn():
+      wrapped_fn(0.5)
+
+    self.assert_trace_line_count(fn, count=15, filtering_enabled=True)
+    self.assert_trace_line_count(fn, count=30, filtering_enabled=False)
+
+  def test_eager_argmax(self):
+    def fn():
+      _ = math_ops.argmax([0, 1], axis=2)
+
+    self.assert_trace_line_count(fn, count=15, filtering_enabled=True)
+    self.assert_trace_line_count(fn, count=30, filtering_enabled=False)
+
+  def test_tfn_argmax(self):
+    @def_function.function
+    def wrapped_fn(x):
+      return math_ops.argmax(x, axis=2)
+
+    def fn():
+      wrapped_fn([0, 1])
+
+    if sys.version_info >= (3, 13):
+      self.assert_trace_line_count(fn, count=16, filtering_enabled=True)
+    else:
+      self.assert_trace_line_count(fn, count=15, filtering_enabled=True)
+    self.assert_trace_line_count(fn, count=25, filtering_enabled=False)
+
+  def test_variable_constructor(self):
+    def fn():
+      _ = variables.Variable()
+
+    self.assert_trace_line_count(fn, count=15, filtering_enabled=True)
+    self.assert_trace_line_count(fn, count=30, filtering_enabled=False)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/util/transform_graph_wrapper.cc

@@ -0,0 +1,78 @@
+/* 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 <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+#include "tensorflow/tools/graph_transforms/transform_graph.h"
+
+namespace py = pybind11;
+
+namespace tensorflow {
+
+std::string TransformGraphWithStringInputs(std::string graph_def_string,
+                                           std::string inputs_string,
+                                           std::string outputs_string,
+                                           std::string transforms_string) {
+  GraphDef graph_def;
+  if (!graph_def.ParseFromString(graph_def_string)) {
+    MaybeRaiseFromStatus(
+        absl::InvalidArgumentError("Couldn't interpret input as a GraphDef"));
+  }
+
+  graph_transforms::TransformParameters params_list;
+  absl::Status parse_status = graph_transforms::ParseTransformParameters(
+      transforms_string, &params_list);
+  if (!parse_status.ok()) {
+    MaybeRaiseFromStatus(parse_status);
+  }
+  std::vector<std::string> inputs = str_util::Split(inputs_string, ',');
+  std::vector<std::string> outputs = str_util::Split(outputs_string, ',');
+
+  absl::Status transform_status = graph_transforms::TransformGraph(
+      inputs, outputs, params_list, &graph_def);
+  if (!transform_status.ok()) {
+    MaybeRaiseFromStatus(transform_status);
+  }
+  std::string result;
+  if (!graph_def.SerializeToString(&result)) {
+    MaybeRaiseFromStatus(
+        absl::InvalidArgumentError("Couldn't serialize output as a GraphDef"));
+  }
+  return result;
+}
+
+}  // namespace tensorflow
+
+PYBIND11_MODULE(_pywrap_transform_graph, m, pybind11::mod_gil_not_used()) {
+  m.def(
+      "TransformGraphWithStringInputs",
+      [](const py::object graph_def_string, const py::object inputs_string,
+         const py::object outputs_string, const py::object transforms_string) {
+        return py::bytes(tensorflow::TransformGraphWithStringInputs(
+            graph_def_string.cast<std::string>(),
+            inputs_string.cast<std::string>(),
+            outputs_string.cast<std::string>(),
+            transforms_string.cast<std::string>()));
+      });
+};

tensorflow/python/util/type_annotations.py

@@ -0,0 +1,59 @@
+# 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.
+# ==============================================================================
+"""Utilities for accessing Python generic type annotations (typing.*)."""
+
+import collections.abc
+import typing
+
+
+def is_generic_union(tp):
+  """Returns true if `tp` is a parameterized typing.Union value."""
+  return (tp is not typing.Union and
+          getattr(tp, '__origin__', None) is typing.Union)
+
+
+def is_generic_tuple(tp):
+  """Returns true if `tp` is a parameterized typing.Tuple value."""
+  return (tp not in (tuple, typing.Tuple) and
+          getattr(tp, '__origin__', None) in (tuple, typing.Tuple))
+
+
+def is_generic_list(tp):
+  """Returns true if `tp` is a parameterized typing.List value."""
+  return (tp not in (list, typing.List) and
+          getattr(tp, '__origin__', None) in (list, typing.List))
+
+
+def is_generic_mapping(tp):
+  """Returns true if `tp` is a parameterized typing.Mapping value."""
+  return (tp not in (collections.abc.Mapping, typing.Mapping) and getattr(
+      tp, '__origin__', None) in (collections.abc.Mapping, typing.Mapping))
+
+
+def is_forward_ref(tp):
+  """Returns true if `tp` is a typing forward reference."""
+  if hasattr(typing, 'ForwardRef'):
+    return isinstance(tp, typing.ForwardRef)
+  elif hasattr(typing, '_ForwardRef'):
+    return isinstance(tp, typing._ForwardRef)  # pylint: disable=protected-access
+  else:
+    return False
+
+
+# Note: typing.get_args was added in Python 3.8.
+if hasattr(typing, 'get_args'):
+  get_generic_type_args = typing.get_args
+else:
+  get_generic_type_args = lambda tp: tp.__args__

tensorflow/python/util/type_annotations_test.py

@@ -0,0 +1,73 @@
+# 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 tensorflow/python/util/type_annotations.py."""
+
+import typing
+from absl.testing import parameterized
+
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
+from tensorflow.python.util import type_annotations
+
+
+class TypeAnnotationsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      (typing.Union[int, float], 'Union'),
+      (typing.Tuple[int, ...], 'Tuple'),
+      (typing.Tuple[int, float, float], 'Tuple'),
+      (typing.Mapping[int, float], 'Mapping'),
+      (typing.Union[typing.Tuple[int], typing.Tuple[int, ...]], 'Union'),
+      # These predicates return False for Generic types w/ no parameters:
+      (typing.Union, None),
+      (typing.Tuple, None),
+      (typing.Mapping, None),
+      (int, None),
+      (12, None),
+  ])
+  def testGenericTypePredicates(self, tp, expected):
+    self.assertEqual(
+        type_annotations.is_generic_union(tp), expected == 'Union')
+    self.assertEqual(
+        type_annotations.is_generic_tuple(tp), expected == 'Tuple')
+    self.assertEqual(
+        type_annotations.is_generic_mapping(tp), expected == 'Mapping')
+
+  @parameterized.parameters([
+      (typing.Union[int, float], (int, float)),
+      (typing.Tuple[int, ...], (int, Ellipsis)),
+      (typing.Tuple[int, float, float], (
+          int,
+          float,
+          float,
+      )),
+      (typing.Mapping[int, float], (int, float)),
+      (typing.Union[typing.Tuple[int],
+                    typing.Tuple[int,
+                                 ...]], (typing.Tuple[int], typing.Tuple[int,
+                                                                         ...])),
+  ])
+  def testGetGenericTypeArgs(self, tp, expected):
+    self.assertEqual(type_annotations.get_generic_type_args(tp), expected)
+
+  def testIsForwardRef(self):
+    tp = typing.Union['B', int]
+    tp_args = type_annotations.get_generic_type_args(tp)
+    self.assertTrue(type_annotations.is_forward_ref(tp_args[0]))
+    self.assertFalse(type_annotations.is_forward_ref(tp_args[1]))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/util/util.h

@@ -0,0 +1,254 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Functions for getting information about kernels registered in the binary.
+#ifndef TENSORFLOW_PYTHON_UTIL_UTIL_H_
+#define TENSORFLOW_PYTHON_UTIL_UTIL_H_
+
+#include <Python.h>
+
+#include <string>
+
+namespace tensorflow {
+namespace swig {
+
+// Implements `tensorflow.util.nest.is_nested`.
+bool IsNested(PyObject* o);
+
+// Implements `tensorflow.util.nest.is_nested_or_composite`.
+bool IsNestedOrComposite(PyObject* o);
+
+// Returns a true if its input is a CompositeTensor or a TypeSpec.
+//
+// Args:
+//   o: the object to check.
+//
+// Returns:
+//   True if the object is a CompositeTensor.
+bool IsCompositeTensor(PyObject* o);
+
+// Returns a true if its input is a TypeSpec, but is not a TensorSpec.
+//
+// Args:
+//   o: the object to check.
+//
+// Returns:
+//   True if the object is a TypeSpec, but is not a TensorSpec.
+bool IsTypeSpec(PyObject* o);
+
+// Implements the same interface as tensorflow.util.nest.is_namedtuple
+// Returns Py_True iff `instance` should be considered a `namedtuple`.
+//
+// Args:
+//   instance: An instance of a Python object.
+//   strict: If True, `instance` is considered to be a `namedtuple` only if
+//       it is a "plain" namedtuple. For instance, a class inheriting
+//       from a `namedtuple` will be considered to be a `namedtuple`
+//       iff `strict=False`.
+//
+// Returns:
+//   True if `instance` is a `namedtuple`.
+PyObject* IsNamedtuple(PyObject* o, bool strict);
+
+// Returns a true if its input is a collections.Mapping.
+//
+// Args:
+//   o: the object to be checked.
+//
+// Returns:
+//   True if the object subclasses mapping.
+bool IsMapping(PyObject* o);
+
+// Returns a true if its input is a collections.MutableMapping.
+//
+// Args:
+//   o: the object to be checked.
+//
+// Returns:
+//   True if the object subclasses mapping.
+bool IsMutableMapping(PyObject* o);
+
+// Returns a true if its input is a (possibly wrapped) tuple.
+//
+// Args:
+//   o: the object to be checked.
+//
+// Returns:
+//   True if the object is a tuple.
+bool IsTuple(PyObject* o);
+
+// Returns a true if its input is a collections.MappingView.
+//
+// Args:
+//   o: the object to be checked.
+//
+// Returns:
+//   True if the object subclasses mapping.
+bool IsMappingView(PyObject* o);
+
+// Returns a true if its input has a `__tf_dispatch__` attribute.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if `o` has a `__tf_dispatch__` attribute.
+bool IsDispatchable(PyObject* o);
+
+// A version of PyMapping_Keys that works in C++11
+//
+// Args:
+//   o: The input to extract keys from
+//
+// Returns:
+//   A new reference to a list of keys in the mapping.
+PyObject* MappingKeys(PyObject* o);
+
+// Returns a true if its input is an instance of an attr.s decorated class.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is an instance of an attr.s decorated class.
+bool IsAttrs(PyObject* o);
+
+// Returns a true if its input is an ops.Tensor.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is a tensor.
+bool IsTensor(PyObject* o);
+
+// Returns true if its input is a tf.TensorSpec.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is a TensorSpec.
+bool IsTensorSpec(PyObject* o);
+
+// Returns a true if its input is an eager.EagerTensor.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is an eager tensor (or mimicking as one).
+bool IsEagerTensorSlow(PyObject* o);
+
+// Returns a true if its input subclasses TensorProtocol.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object implements TensorProtocol.
+bool IsTensorProtocol(PyObject* o);
+
+// Returns a true if its input is a core.Value type.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is a core.Value type.
+bool IsCoreTypeValue(PyObject* o);
+
+// Returns a true if its input is a ResourceVariable.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is a ResourceVariable.
+bool IsResourceVariable(PyObject* o);
+
+// Returns a true if its input is an OwnedIterator.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is an OwnedIterator.
+bool IsOwnedIterator(PyObject* o);
+
+// Returns a true if its input is a Variable.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is a Variable.
+bool IsVariable(PyObject* o);
+
+// Returns a true if its input is an ops.IndexesSlices.
+//
+// Args:
+//   o: the input to be checked.
+//
+// Returns:
+//   True if the object is an ops.IndexedSlices.
+bool IsIndexedSlices(PyObject* o);
+
+// Implements the same interface as tensorflow.util.nest.same_namedtuples
+// Returns Py_True iff the two namedtuples have the same name and fields.
+// Raises RuntimeError if `o1` or `o2` don't look like namedtuples (don't have
+// '_fields' attribute).
+PyObject* SameNamedtuples(PyObject* o1, PyObject* o2);
+
+// Implements `tensorflow.util.nest.assert_same_structrure`.
+PyObject* AssertSameStructure(PyObject* o1, PyObject* o2, bool check_types,
+                              bool expand_composites);
+
+// Implements `tensorflow.util.nest.flatten`.
+PyObject* Flatten(PyObject* nested, bool expand_composites = false);
+
+// The tensorflow.python.data package has its own nest utility that follows very
+// slightly different semantics for its functions than the tensorflow.python
+// nest utility. Returns True if its input is a nested structure for tf.data.
+//
+// Main differences are (this is copied from nest.py in the
+// tensorflow.data.util):
+//
+// 1. It removes support for lists as a level of nesting in nested structures.
+// 2. It adds support for `SparseTensorValue` as an atomic element.
+bool IsNestedForData(PyObject* o);
+
+// Flatten specialized for `tf.data`. Additional comments about
+// difference in functionality can be found in nest.py in
+// `tensorflow.python.data.util` and in the comments for Flatten above.
+PyObject* FlattenForData(PyObject* nested);
+
+// AssertSameStructure specialized for `tf.data`.
+PyObject* AssertSameStructureForData(PyObject* o1, PyObject* o2,
+                                     bool check_types);
+
+// Registers a Python object so it can be looked up from c++.  The set of
+// valid names, and the expected values for those names, are listed in
+// the documentation for `RegisteredPyObjects`.  Returns PyNone.
+PyObject* RegisterPyObject(PyObject* name, PyObject* value);
+
+// Returns a borrowed reference to an object that was registered with
+// RegisterPyObject.  (Do not call Py_DECREF on the result).
+PyObject* GetRegisteredPyObject(const std::string& name);
+
+}  // namespace swig
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_UTIL_UTIL_H_

tensorflow/python/util/util_wrapper.cc

@@ -0,0 +1,333 @@
+/* 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.
+==============================================================================*/
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/util/util.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/util/util.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_utils, m, pybind11::mod_gil_not_used()) {
+  m.doc() = R"pbdoc(
+    _pywrap_utils
+    -----
+  )pbdoc";
+  m.def("RegisterPyObject", [](const py::handle& name, const py::handle& type) {
+    return tensorflow::PyoOrThrow(
+        tensorflow::swig::RegisterPyObject(name.ptr(), type.ptr()));
+  });
+  m.def(
+      "IsTensor",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsTensor(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Check if an object is a Tensor.
+    )pbdoc");
+  m.def(
+      "IsNested",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsNested(o.ptr());
+        return result;
+      },
+      R"pbdoc(
+      Refer to `tf.nest.is_nested`.
+    )pbdoc");
+  m.def(
+      "IsNestedOrComposite",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsNestedOrComposite(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns true if its input is a sequence or a `CompositeTensor`.
+
+      Args:
+        seq: an input sequence.
+
+      Returns:
+        True if the sequence is a not a string and is a collections.Sequence or a
+        dict or a CompositeTensor or a TypeSpec (except string and TensorSpec).
+    )pbdoc");
+  m.def(
+      "IsCompositeTensor",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsCompositeTensor(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns true if its input is a `CompositeTensor`.
+
+      Args:
+        seq: an input sequence.
+
+      Returns:
+        True if the sequence is a CompositeTensor.
+    )pbdoc");
+  m.def(
+      "IsTypeSpec",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsTypeSpec(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns true if its input is a `TypeSpec`, but is not a `TensorSpec`.
+
+      Args:
+        seq: an input sequence.
+
+      Returns:
+        True if the sequence is a `TypeSpec`, but is not a `TensorSpec`.
+    )pbdoc");
+  m.def(
+      "IsNamedtuple",
+      [](const py::handle& o, bool strict) {
+        return tensorflow::PyoOrThrow(
+            tensorflow::swig::IsNamedtuple(o.ptr(), strict));
+      },
+      R"pbdoc(
+      Check if an object is a NamedTuple.
+    )pbdoc");
+  m.def(
+      "IsMapping",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsMapping(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns True if `instance` is a `collections.Mapping`.
+
+      Args:
+        instance: An instance of a Python object.
+
+      Returns:
+        True if `instance` is a `collections.Mapping`.
+    )pbdoc");
+  m.def(
+      "IsMutableMapping",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsMutableMapping(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns True if `instance` is a `collections.MutableMapping`.
+
+      Args:
+        instance: An instance of a Python object.
+
+      Returns:
+        True if `instance` is a `collections.MutableMapping`.
+    )pbdoc");
+  m.def(
+      "IsMappingView",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsMappingView(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns True if considered a mapping view for the purposes of Flatten()`.
+
+      Args:
+        instance: An instance of a Python object.
+
+      Returns:
+        True if considered a mapping view for the purposes of Flatten().
+    )pbdoc");
+  m.def(
+      "IsAttrs",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsAttrs(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns True if `instance` is an instance of an `attr.s` decorated class.
+
+      Args:
+        instance: An instance of a Python object.
+
+      Returns:
+        True if `instance` is an instance of an `attr.s` decorated class.
+    )pbdoc");
+  m.def(
+      "SameNamedtuples",
+      [](const py::handle& o1, const py::handle& o2) {
+        return tensorflow::PyoOrThrow(
+            tensorflow::swig::SameNamedtuples(o1.ptr(), o2.ptr()));
+      },
+      R"pbdoc(
+      Returns True if the two namedtuples have the same name and fields.
+    )pbdoc");
+  m.def(
+      "AssertSameStructure",
+      [](const py::handle& o1, const py::handle& o2, bool check_types,
+         bool expand_composites) {
+        bool result = tensorflow::swig::AssertSameStructure(
+            o1.ptr(), o2.ptr(), check_types, expand_composites);
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns True if the two structures are nested in the same way.
+    )pbdoc");
+  m.def(
+      "Flatten",
+      [](const py::handle& o, bool expand_composites) {
+        return tensorflow::PyoOrThrow(
+            tensorflow::swig::Flatten(o.ptr(), expand_composites));
+      },
+      R"pbdoc(
+      Refer to `tf.nest.flatten`.
+    )pbdoc");
+  m.def(
+      "IsNestedForData",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsNestedForData(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns a true if `seq` is a nested structure for tf.data.
+
+      NOTE(mrry): This differs from `tensorflow.python.util.nest.is_nested()`,
+      which *does* treat a Python list as a sequence. For ergonomic
+      reasons, `tf.data` users would prefer to treat lists as
+      implicit `tf.Tensor` objects, and dicts as (nested) sequences.
+
+      Args:
+        seq: an input sequence.
+
+      Returns:
+        True if the sequence is a not a string or list and is a
+        collections.Sequence.
+    )pbdoc");
+  m.def(
+      "FlattenForData",
+      [](const py::handle& o) {
+        return tensorflow::PyoOrThrow(
+            tensorflow::swig::FlattenForData(o.ptr()));
+      },
+      R"pbdoc(
+      Returns a flat sequence from a given nested structure.
+
+      If `nest` is not a sequence, this returns a single-element list: `[nest]`.
+
+      Args:
+        nest: an arbitrarily nested structure or a scalar object.
+          Note, numpy arrays are considered scalars.
+
+      Returns:
+        A Python list, the flattened version of the input.
+    )pbdoc");
+  m.def(
+      "AssertSameStructureForData",
+      [](const py::handle& o1, const py::handle& o2, bool check_types) {
+        bool result = tensorflow::swig::AssertSameStructureForData(
+            o1.ptr(), o2.ptr(), check_types);
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns True if the two structures are nested in the same way in particular tf.data.
+    )pbdoc");
+  m.def(
+      "IsResourceVariable",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsResourceVariable(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns 1 if `o` is a ResourceVariable.
+
+      Args:
+        instance: An instance of a Python object.
+
+      Returns:
+        True if `instance` is a `ResourceVariable`.
+    )pbdoc");
+  m.def(
+      "IsVariable",
+      [](const py::handle& o) {
+        bool result = tensorflow::swig::IsVariable(o.ptr());
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns 1 if `o` is a Variable.
+
+      Args:
+        instance: An instance of a Python object.
+
+      Returns:
+        True if `instance` is a `Variable`.
+    )pbdoc");
+  m.def(
+      "IsDataTypeSupportedByOneDNNOnThisCPU",
+      [](const tensorflow::DataType& dt) {
+        bool result = tensorflow::IsDataTypeSupportedByOneDNNOnThisCPU(dt);
+        if (PyErr_Occurred()) {
+          throw py::error_already_set();
+        }
+        return result;
+      },
+      R"pbdoc(
+      Returns true if input type is supported on CPU when oneDNN is enabled
+
+      Args:
+       type: an input data type
+
+      Returns:
+        True if input type is supported on CPU when oneDNN is enabled.
+        False otherwise.
+    )pbdoc");
+}