tensorflow/python/framework/test_util.py

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

# pylint: disable=invalid-name
"""Test utils for tensorflow."""

import collections
from collections import OrderedDict
from collections.abc import Iterable, Iterator, Callable, Collection, Sequence
import contextlib
import functools
import gc
import itertools
import math
import os
import random
import re
import tempfile
import threading
import time
from typing import Any, cast, Union, Optional, overload, TypeVar
import unittest

from absl.testing import parameterized
import numpy as np

from google.protobuf import descriptor_pool
from google.protobuf import message
from google.protobuf import text_format
from tensorflow.core.config import flags
from tensorflow.core.framework import graph_pb2
from tensorflow.core.framework import node_def_pb2
from tensorflow.core.framework import tensor_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python import pywrap_sanitizers
from tensorflow.python import tf2
from tensorflow.python.client import device_lib
from tensorflow.python.client import pywrap_tf_session
from tensorflow.python.client import session as s
from tensorflow.python.compat import v2_compat
from tensorflow.python.compat.compat import forward_compatibility_horizon
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import _test_metrics_util
from tensorflow.python.framework import config
from tensorflow.python.framework import device as pydev
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import gpu_util
from tensorflow.python.framework import importer
from tensorflow.python.framework import indexed_slices
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.framework import tfrt_utils
from tensorflow.python.framework import versions
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import control_flow_util_v2
from tensorflow.python.ops import gen_sparse_ops
from tensorflow.python.ops import gen_sync_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import script_ops
from tensorflow.python.ops import summary_ops_v2
from tensorflow.python.ops import variables


from tensorflow.python.ops.ragged import ragged_ops  # pylint: disable=unused-import
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged import ragged_tensor_value
from tensorflow.python.platform import _pywrap_stacktrace_handler
from tensorflow.python.platform import googletest
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
from tensorflow.python.util import _pywrap_util_port
from tensorflow.python.util import compat
from tensorflow.python.util import deprecation
from tensorflow.python.util import nest
from tensorflow.python.util import tf_decorator
from tensorflow.python.util import tf_inspect
from tensorflow.python.util import traceback_utils
from tensorflow.python.util.compat import collections_abc
from tensorflow.python.util.protobuf import compare
from tensorflow.python.util.tf_export import tf_export


_F = TypeVar("_F", bound=Callable[..., Any])
_T = TypeVar("_T")
_TC = TypeVar("_TC", bound=type["TensorFlowTestCase"])


# If the below import is made available through the BUILD rule, then this
# function is overridden and will instead return True and cause Tensorflow
# graphs to be compiled with XLA.
def is_xla_enabled() -> bool:
  return False


# pytype: disable=import-error
try:
  from tensorflow.python.framework.is_xla_test_true import is_xla_enabled  # pylint: disable=g-import-not-at-top, unused-import
except Exception:  # pylint: disable=broad-except
  pass


# Uses the same mechanism as above to selectively enable/disable MLIR
# compilation.
def is_mlir_bridge_enabled() -> Optional[bool]:
  return None


try:
  from tensorflow.python.framework.is_mlir_bridge_test_false import is_mlir_bridge_enabled  # pylint: disable=g-import-not-at-top, unused-import
except ImportError:
  try:
    from tensorflow.python.framework.is_mlir_bridge_test_true import is_mlir_bridge_enabled  # pylint: disable=g-import-not-at-top, unused-import
  except ImportError:
    pass
# pytype: enable=import-error


def is_asan_enabled() -> bool:
  """Check if ASAN is enabled."""
  return pywrap_sanitizers.is_asan_enabled()


def is_msan_enabled() -> bool:
  """Check if MSAN is enabled."""
  return pywrap_sanitizers.is_msan_enabled()


def is_tsan_enabled() -> bool:
  """Check if TSAN is enabled."""
  return pywrap_sanitizers.is_tsan_enabled()


def is_ubsan_enabled() -> bool:
  """Check if UBSAN is enabled."""
  return pywrap_sanitizers.is_ubsan_enabled()


def _get_object_count_by_type(
    exclude: Iterable[Any] = (),
) -> collections.Counter[str]:
  return (
      collections.Counter([type(obj).__name__ for obj in gc.get_objects()]) -
      collections.Counter([type(obj).__name__ for obj in exclude]))


@tf_export("test.gpu_device_name")
def gpu_device_name() -> str:
  """Returns the name of a GPU device if available or a empty string.

  This method should only be used in tests written with `tf.test.TestCase`.

  >>> class MyTest(tf.test.TestCase):
  ...
  ...   def test_add_on_gpu(self):
  ...     if not tf.test.is_built_with_gpu_support():
  ...       self.skipTest("test is only applicable on GPU")
  ...
  ...     with tf.device(tf.test.gpu_device_name()):
  ...       self.assertEqual(tf.math.add(1.0, 2.0), 3.0)

  """
  for x in device_lib.list_local_devices():
    if x.device_type == "GPU":
      return compat.as_str(x.name)
  return ""


def assert_ops_in_graph(
    expected_ops: dict[str, str], graph: ops.Graph
) -> dict[str, node_def_pb2.NodeDef]:
  """Assert all expected operations are found.

  Args:
    expected_ops: `dict<string, string>` of op name to op type.
    graph: Graph to check.

  Returns:
    `dict<string, node>` of node name to node.

  Raises:
    ValueError: If the expected ops are not present in the graph.
  """
  actual_ops: dict[str, node_def_pb2.NodeDef] = {}
  gd = cast(graph_pb2.GraphDef, graph.as_graph_def())
  for node in gd.node:
    if node.name in expected_ops:
      if expected_ops[node.name] != node.op:
        raise ValueError("Expected op for node %s is different. %s vs %s" %
                         (node.name, expected_ops[node.name], node.op))
      actual_ops[node.name] = node
  if set(expected_ops.keys()) != set(actual_ops.keys()):
    raise ValueError("Not all expected ops are present. Expected %s, found %s" %
                     (expected_ops.keys(), actual_ops.keys()))
  return actual_ops


@tf_export("test.assert_equal_graph_def", v1=[])
def assert_equal_graph_def_v2(
    expected: graph_pb2.GraphDef, actual: graph_pb2.GraphDef
) -> None:
  """Asserts that two `GraphDef`s are (mostly) the same.

  Compares two `GraphDef` protos for equality, ignoring versions and ordering of
  nodes, attrs, and control inputs.  Node names are used to match up nodes
  between the graphs, so the naming of nodes must be consistent. This function
  ignores randomized attribute values that may appear in V2 checkpoints.

  Args:
    expected: The `GraphDef` we expected.
    actual: The `GraphDef` we have.

  Raises:
    AssertionError: If the `GraphDef`s do not match.
    TypeError: If either argument is not a `GraphDef`.
  """
  assert_equal_graph_def(actual, expected, checkpoint_v2=True,
                         hash_table_shared_name=True)


@tf_export(v1=["test.assert_equal_graph_def"])
def assert_equal_graph_def_v1(
    actual: graph_pb2.GraphDef,
    expected: graph_pb2.GraphDef,
    checkpoint_v2: bool = False,
    hash_table_shared_name: bool = False
) -> None:
  """Asserts that two `GraphDef`s are (mostly) the same.

  Compares two `GraphDef` protos for equality, ignoring versions and ordering of
  nodes, attrs, and control inputs.  Node names are used to match up nodes
  between the graphs, so the naming of nodes must be consistent.

  Args:
    actual: The `GraphDef` we have.
    expected: The `GraphDef` we expected.
    checkpoint_v2: boolean determining whether to ignore randomized attribute
      values that appear in V2 checkpoints.
    hash_table_shared_name: boolean determining whether to ignore randomized
      shared_names that appear in HashTableV2 op defs.

  Raises:
    AssertionError: If the `GraphDef`s do not match.
    TypeError: If either argument is not a `GraphDef`.
  """
  assert_equal_graph_def(actual, expected, checkpoint_v2,
                         hash_table_shared_name)


def assert_equal_graph_def(
    actual: graph_pb2.GraphDef,
    expected: graph_pb2.GraphDef,
    checkpoint_v2: bool = False,
    hash_table_shared_name: bool = False
)-> None:
  if not isinstance(actual, graph_pb2.GraphDef):
    raise TypeError("Expected tf.GraphDef for actual, got %s" %
                    type(actual).__name__)
  if not isinstance(expected, graph_pb2.GraphDef):
    raise TypeError("Expected tf.GraphDef for expected, got %s" %
                    type(expected).__name__)

  if checkpoint_v2:
    _strip_checkpoint_v2_randomized(actual)
    _strip_checkpoint_v2_randomized(expected)

  if hash_table_shared_name:
    _strip_hash_table_shared_name(actual)
    _strip_hash_table_shared_name(expected)

  diff = pywrap_tf_session.EqualGraphDefWrapper(actual.SerializeToString(),
                                                expected.SerializeToString())
  if diff:
    raise AssertionError(compat.as_str(diff))


def assert_meta_graph_protos_equal(
    tester: "TensorFlowTestCase",
    a: meta_graph_pb2.MetaGraphDef,
    b: meta_graph_pb2.MetaGraphDef,
) -> None:
  """Compares MetaGraphDefs `a` and `b` in unit test class `tester`."""
  # Carefully check the collection_defs
  tester.assertEqual(set(a.collection_def), set(b.collection_def))
  collection_keys = a.collection_def.keys()
  for k in collection_keys:
    a_value = a.collection_def[k]
    b_value = b.collection_def[k]
    proto_type = cast(type[message.Message], ops.get_collection_proto_type(k))
    if proto_type:
      a_proto = proto_type()
      b_proto = proto_type()
      # Number of entries in the collections is the same
      tester.assertEqual(
          len(a_value.bytes_list.value), len(b_value.bytes_list.value))
      for (a_value_item, b_value_item) in zip(a_value.bytes_list.value,
                                              b_value.bytes_list.value):
        a_proto.ParseFromString(a_value_item)
        b_proto.ParseFromString(b_value_item)
        tester.assertProtoEquals(a_proto, b_proto)
    else:
      tester.assertEqual(a_value, b_value)
  # Compared the fields directly, remove their raw values from the
  # proto comparison below.
  a.ClearField("collection_def")
  b.ClearField("collection_def")

  # Check the graph_defs.
  assert_equal_graph_def(a.graph_def, b.graph_def, checkpoint_v2=True)
  # Check graph_def versions (ignored by assert_equal_graph_def).
  tester.assertProtoEquals(a.graph_def.versions, b.graph_def.versions)
  # Compared the fields directly, remove their raw values from the
  # proto comparison below.
  a.ClearField("graph_def")
  b.ClearField("graph_def")

  tester.assertProtoEquals(a, b)


# Matches attributes named via _SHARDED_SUFFIX in
# tensorflow/python/training/saver.py
_SHARDED_SAVE_OP_PATTERN = "_temp_[0-9a-z]{32}/part"


def _strip_checkpoint_v2_randomized(graph_def: graph_pb2.GraphDef) -> None:
  for node in graph_def.node:
    delete_keys: list[str] = []
    for attr_key in node.attr:
      attr_tensor_value = cast(
          tensor_pb2.TensorProto, node.attr[attr_key].tensor)
      if attr_tensor_value and len(attr_tensor_value.string_val) == 1:
        attr_tensor_string_value = attr_tensor_value.string_val[0]
        if (attr_tensor_string_value and
            re.match(compat.as_bytes(_SHARDED_SAVE_OP_PATTERN),
                     attr_tensor_string_value)):
          delete_keys.append(attr_key)
    for attr_key in delete_keys:
      del node.attr[attr_key]


_TABLE_SHARED_NAME_PATTERN = r"hash_table_[0-9a-z\-]+"


def _strip_hash_table_shared_name(graph_def: graph_pb2.GraphDef) -> None:
  for node in graph_def.node:
    delete_keys: list[str] = []
    if node.op == "HashTableV2" and "shared_name" in node.attr:
      if re.match(compat.as_bytes(_TABLE_SHARED_NAME_PATTERN),
                  node.attr["shared_name"].s):
        delete_keys.append("shared_name")
    for attr_key in delete_keys:
      del node.attr[attr_key]


def IsGoogleCudaEnabled() -> bool:
  return _pywrap_util_port.IsGoogleCudaEnabled()


def IsBuiltWithROCm() -> bool:
  return _pywrap_util_port.IsBuiltWithROCm()


def IsBuiltWithXLA() -> bool:
  return _pywrap_util_port.IsBuiltWithXLA()


def IsBuiltWithNvcc() -> bool:
  return _pywrap_util_port.IsBuiltWithNvcc()


def IsCPUTargetAvailable(target):
  if target.startswith("arm"):
    return _pywrap_util_port.IsAArch32Available()
  elif target.startswith("aarch64"):
    return _pywrap_util_port.IsAArch64Available()
  elif target.startswith("ppc") or target.startswith("powerpc"):
    return _pywrap_util_port.IsPowerPCAvailable()
  elif target.startswith("x86"):
    return _pywrap_util_port.IsX86Available()
  elif target.startswith("s390x"):
    return _pywrap_util_port.IsSystemZAvailable()
  assert False, "Unknown CPU target: " + target


def GpuSupportsHalfMatMulAndConv() -> bool:
  return _pywrap_util_port.GpuSupportsHalfMatMulAndConv()


def IsMklEnabled() -> bool:
  return _pywrap_util_port.IsMklEnabled()


def InstallStackTraceHandler() -> None:
  _pywrap_stacktrace_handler.InstallStacktraceHandler()


def NHWCToNCHW(
    input_tensor: Union[tensor_lib.Tensor, list[int]]
) -> Union[tensor_lib.Tensor, list[int]]:
  """Converts the input from the NHWC format to NCHW.

  Args:
    input_tensor: a 3-, 4-, or 5-D tensor, or an array representing shape

  Returns:
    converted tensor or shape array
  """
  # tensor dim -> new axis order
  new_axes = {3: [0, 2, 1], 4: [0, 3, 1, 2], 5: [0, 4, 1, 2, 3]}
  if isinstance(input_tensor, tensor_lib.Tensor):
    ndims = input_tensor.shape.ndims
    return array_ops.transpose(input_tensor, new_axes[ndims])
  else:
    ndims = len(input_tensor)
    return [input_tensor[a] for a in new_axes[ndims]]


def NHWCToNCHW_VECT_C(
    input_shape_or_tensor: Union[tensor_lib.Tensor, list[int]]
)-> Union[tensor_lib.Tensor, list[int]]:
  """Transforms the input from the NHWC layout to NCHW_VECT_C layout.

  Note: Does not include quantization or type conversion steps, which should
  be applied afterwards.

  Args:
    input_shape_or_tensor: a 4- or 5-D tensor, or an array representing shape

  Returns:
    tensor or shape array transformed into NCHW_VECT_C

  Raises:
    ValueError: if last dimension of `input_shape_or_tensor` is not evenly
        divisible by 4.
  """
  permutations = {5: [0, 3, 1, 2, 4], 6: [0, 4, 1, 2, 3, 5]}
  is_tensor = isinstance(input_shape_or_tensor, tensor_lib.Tensor)
  temp_shape: list[int] = (
      input_shape_or_tensor.shape.as_list()
      if is_tensor else input_shape_or_tensor)
  if temp_shape[-1] % 4 != 0:
    raise ValueError(
        "Last dimension of input must be evenly divisible by 4 to convert to "
        "NCHW_VECT_C.")
  temp_shape[-1] //= 4
  temp_shape.append(4)
  permutation = permutations[len(temp_shape)]
  if is_tensor:
    t = array_ops.reshape(input_shape_or_tensor, temp_shape)
    return array_ops.transpose(t, permutation)
  else:
    return [temp_shape[a] for a in permutation]


def NCHW_VECT_CToNHWC(
    input_shape_or_tensor: Union[tensor_lib.Tensor, list[int]]
) -> Union[tensor_lib.Tensor, list[int]]:
  """Transforms the input from the NCHW_VECT_C layout to NHWC layout.

  Note: Does not include de-quantization or type conversion steps, which should
  be applied beforehand.

  Args:
    input_shape_or_tensor: a 5- or 6-D tensor, or an array representing shape

  Returns:
    tensor or shape array transformed into NHWC

  Raises:
    ValueError: if last dimension of `input_shape_or_tensor` is not 4.
  """
  permutations = {5: [0, 2, 3, 1, 4], 6: [0, 2, 3, 4, 1, 5]}
  is_tensor = isinstance(input_shape_or_tensor, tensor_lib.Tensor)
  input_shape: list[int] = (
      input_shape_or_tensor.shape.as_list()
      if is_tensor else input_shape_or_tensor)
  if input_shape[-1] != 4:
    raise ValueError("Last dimension of NCHW_VECT_C must be 4.")
  permutation = permutations[len(input_shape)]
  nhwc_shape = [input_shape[a] for a in permutation[:-1]]
  nhwc_shape[-1] *= input_shape[-1]
  if is_tensor:
    t = array_ops.transpose(input_shape_or_tensor, permutation)
    return array_ops.reshape(t, nhwc_shape)
  else:
    return nhwc_shape


def NCHWToNHWC(
    input_tensor: Union[tensor_lib.Tensor, list[int]]
) -> Union[tensor_lib.Tensor, list[int]]:
  """Converts the input from the NCHW format to NHWC.

  Args:
    input_tensor: a 4- or 5-D tensor, or an array representing shape

  Returns:
    converted tensor or shape array
  """
  # tensor dim -> new axis order
  new_axes = {4: [0, 2, 3, 1], 5: [0, 2, 3, 4, 1]}
  if isinstance(input_tensor, tensor_lib.Tensor):
    ndims = input_tensor.shape.ndims
    return array_ops.transpose(input_tensor, new_axes[ndims])
  else:
    ndims = len(input_tensor)
    return [input_tensor[a] for a in new_axes[ndims]]


def skip_if(condition: Union[Callable[[], bool], bool]) -> Callable[[_F], _F]:
  """Skips the decorated function if condition is or evaluates to True.

  Args:
    condition: Either an expression that can be used in "if not condition"
      statement, or a callable whose result should be a boolean.

  Returns:
    The wrapped function
  """

  def real_skip_if(fn: _F) -> _F:

    def wrapper(*args, **kwargs):
      if callable(condition):
        skip = condition()
      else:
        skip = condition
      if not skip:
        return fn(*args, **kwargs)

    return wrapper

  return real_skip_if


@contextlib.contextmanager
def skip_if_error(
    test_obj: unittest.TestCase,
    error_type: type[Exception],
    messages: Union[str, list[str], None] = None
) -> Iterator[None]:
  """Context manager to skip cases not considered failures by the tests.

  Note that this does not work if used in setUpClass/tearDownClass.
  Usage in setUp/tearDown works fine just like regular test methods.

  Args:
    test_obj: A test object provided as `self` in the test methods; this object
      is usually an instance of `unittest.TestCase`'s subclass and should have
      `skipTest` method.
    error_type: The error type to skip. Note that if `messages` are given, both
      `error_type` and `messages` need to match for the test to be skipped.
    messages: Optional, a string or list of strings. If `None`, the test will be
      skipped if `error_type` matches what is raised; otherwise, the test is
      skipped if any of the `messages` is contained in the message of the error
      raised, and `error_type` matches the error raised.

  Yields:
    Nothing.
  """
  if messages:
    messages = nest.flatten(messages)
  try:
    yield
  except error_type as e:
    if not messages or any(message in str(e) for message in messages):
      test_obj.skipTest("Skipping error: {}: {}".format(type(e), str(e)))
    else:
      raise


def enable_c_shapes(fn: _F) -> _F:
  """No-op. TODO(b/74620627): Remove this."""
  return fn


def with_c_shapes(cls: type[_T]) -> type[_T]:
  """No-op. TODO(b/74620627): Remove this."""
  return cls


def enable_control_flow_v2(fn: _F) -> _F:
  """Decorator for enabling CondV2 and WhileV2 on a test.

  Note this enables using CondV2 and WhileV2 after running the test class's
  setup/teardown methods.

  In addition to this, callers must import the while_v2 module in order to set
  the _while_v2 module in control_flow_ops.

  Args:
    fn: the function to be wrapped

  Returns:
    The wrapped function
  """

  def wrapper(*args, **kwargs):
    enable_control_flow_v2_old = control_flow_util.ENABLE_CONTROL_FLOW_V2
    control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
    try:
      return fn(*args, **kwargs)
    finally:
      control_flow_util.ENABLE_CONTROL_FLOW_V2 = enable_control_flow_v2_old

  return wrapper


def with_control_flow_v2(cls: _TC) -> _TC:
  """Adds methods that call original methods with WhileV2 and CondV2 enabled.

  Note this enables CondV2 and WhileV2 in new methods after running the test
  class's setup method.

  In addition to this, callers must import the while_v2 module in order to set
  the _while_v2 module in control_flow_ops.

  If a test function has _disable_control_flow_v2 attr set to True (using the
  @disable_control_flow_v2 decorator), the v2 function is not generated for it.

  Example:

  @test_util.with_control_flow_v2
  class ControlFlowTest(test.TestCase):

    def testEnabledForV2(self):
      ...

    @test_util.disable_control_flow_v2("b/xyzabc")
    def testDisabledForV2(self):
      ...

  Generated class:
  class ControlFlowTest(test.TestCase):

    def testEnabledForV2(self):
      ...

    def testEnabledForV2WithControlFlowV2(self):
      // Enable V2 flags.
      testEnabledForV2(self)
      // Restore V2 flags.

    def testDisabledForV2(self):
      ...

  Args:
    cls: class to decorate

  Returns:
    cls with new test methods added
  """
  if control_flow_util.ENABLE_CONTROL_FLOW_V2:
    return cls

  for name, value in cls.__dict__.copy().items():
    if (callable(value) and
        name.startswith(unittest.TestLoader.testMethodPrefix) and
        not getattr(value, "_disable_control_flow_v2", False)):
      setattr(cls, name + "WithControlFlowV2", enable_control_flow_v2(value))
  return cls


def disable_control_flow_v2(unused_msg: str) -> Callable[[_F], _F]:
  """Decorator for a function in a with_control_flow_v2 enabled test class.

  Blocks the function from being run with v2 control flow ops.

  Args:
    unused_msg: Reason for disabling.

  Returns:
    The wrapped function with _disable_control_flow_v2 attr set to True.
  """

  def wrapper(func: _F) -> _F:
    func._disable_control_flow_v2 = True
    return func

  return wrapper


def enable_output_all_intermediates(fn: _F) -> _F:
  """Force-enable outputing all intermediates from functional control flow ops.

  Args:
    fn: the function to be wrapped

  Returns:
    The wrapped function
  """

  def wrapper(*args, **kwargs):
    output_all_intermediates_old = \
        control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
    control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = True
    try:
      return fn(*args, **kwargs)
    finally:
      control_flow_util_v2._EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = \
          output_all_intermediates_old

  return wrapper


def assert_no_new_pyobjects_executing_eagerly(
    warmup_iters: int = 2,
) -> Callable[[Callable[..., Any]], Callable[..., None]]:
  """Decorator for asserting that no new Python objects persist after a test.

  Returns a decorator that runs the test multiple times executing eagerly,
  first as a warmup and then to let objects accumulate. The warmup helps ignore
  caches which do not grow as the test is run repeatedly.

  Useful for checking that there are no missing Py_DECREFs in the C exercised by
  a bit of Python.

  Args:
    warmup_iters: The numer of warmup iterations, excluded from measuring.

  Returns:
    A decorator function which can be applied to the test function.
  """

  def wrap_f(f: Callable[..., Any]) -> Callable[..., None]:
    def decorator(self: "TensorFlowTestCase", *args, **kwargs) -> None:
      """Warms up, gets object counts, runs the test, checks for new objects."""
      with context.eager_mode():
        gc.disable()
        # Python 3.11 removed "errors" and "skipped" as members of
        # unittest.case._Outcome so get them from the test result object
        # instead.
        test_errors = None
        test_skipped = None
        if hasattr(self._outcome, "errors"):
          test_errors = self._outcome.errors
          test_skipped = self._outcome.skipped
        else:
          test_errors = self._outcome.result.errors
          test_skipped = self._outcome.result.skipped
        # Run the test 2 times as warmup, in an attempt to fill up caches, which
        # should not grow as the test is run repeatedly below.
        #
        # TODO(b/117156879): Running warmup twice is black magic; we have seen
        # tests that fail with 1 warmup run, and pass with 2, on various
        # versions of python2.7.x.
        for _ in range(warmup_iters):
          f(self, *args, **kwargs)
        # Since we aren't in the normal test lifecycle, we need to manually run
        # cleanups to clear out their object references.
        self.doCleanups()

        # Some objects are newly created by _get_object_count_by_type().  So
        # create and save as a dummy variable to include it as a baseline.
        obj_count_by_type = _get_object_count_by_type()
        gc.collect()

        # Make sure any registered functions are cleaned up in the C++ runtime.
        registered_function_names = context.context().list_function_names()

        # unittest.doCleanups adds to self._outcome with each unwound call.
        # These objects are retained across gc collections so we exclude them
        # from the object count calculation.
        obj_count_by_type = _get_object_count_by_type(
            exclude=gc.get_referents(test_errors, test_skipped))

        if ops.has_default_graph():
          collection_sizes_before = {
              collection: len(ops.get_collection(collection))
              for collection in ops.get_default_graph().collections
          }
        for _ in range(3):
          f(self, *args, **kwargs)
        # Since we aren't in the normal test lifecycle, we need to manually run
        # cleanups to clear out their object references.
        self.doCleanups()
        # Note that gc.get_objects misses anything that isn't subject to garbage
        # collection (C types). Collections are a common source of leaks, so we
        # test for collection sizes explicitly.
        if ops.has_default_graph():
          for collection_key in ops.get_default_graph().collections:
            collection = ops.get_collection(collection_key)
            size_before = collection_sizes_before.get(collection_key, 0)
            if len(collection) > size_before:
              raise AssertionError(
                  ("Collection %s increased in size from "
                   "%d to %d (current items %s).") %
                  (collection_key, size_before, len(collection), collection))
            # Make sure our collection checks don't show up as leaked memory by
            # removing references to temporary variables.
            del collection
            del collection_key
            del size_before
          del collection_sizes_before
        gc.collect()

        # There should be no new Python objects hanging around.
        obj_count_by_type = (
            _get_object_count_by_type(
                exclude=gc.get_referents(test_errors, test_skipped)) -
            obj_count_by_type)

        # There should be no newly registered functions hanging around.
        leftover_functions = (
            context.context().list_function_names() - registered_function_names)
        assert not leftover_functions, (
            "The following functions were newly created: %s" %
            leftover_functions)

        # In some cases (specifically on MacOS), new_count is somehow
        # smaller than previous_count.
        # Using plain assert because not all classes using this decorator
        # have assertLessEqual
        assert not obj_count_by_type, (
            "The following objects were newly created: %s" %
            str(obj_count_by_type))
        gc.enable()
    return tf_decorator.make_decorator(f, decorator)

  return wrap_f


def assert_no_new_tensors(f: _F) -> _F:
  """Decorator for asserting that no new Tensors persist after a test.

  Mainly useful for checking that code using the Python C API has correctly
  manipulated reference counts.

  Clears the caches that it knows about, runs the garbage collector, then checks
  that there are no Tensor or Tensor-like objects still around. This includes
  Tensors to which something still has a reference (e.g. from missing
  Py_DECREFs) and uncollectable cycles (i.e. Python reference cycles where one
  of the objects has __del__ defined).

  Args:
    f: The test case to run.

  Returns:
    The decorated test case.
  """

  def decorator(self: "TensorFlowTestCase", **kwargs):
    """Finds existing Tensors, runs the test, checks for new Tensors."""

    def _is_tensorflow_object(obj) -> bool:
      try:
        return isinstance(obj,
                          (tensor_lib.Tensor, variables.Variable,
                           tensor_shape.Dimension, tensor_shape.TensorShape))
      except (ReferenceError, AttributeError):
        # If the object no longer exists, we don't care about it.
        return False

    tensors_before = set(
        id(obj) for obj in gc.get_objects() if _is_tensorflow_object(obj))
    outside_executed_eagerly = cast(bool, context.executing_eagerly())
    # Run the test in a new graph so that collections get cleared when it's
    # done, but inherit the graph key so optimizers behave.
    outside_graph_key = ops.get_default_graph()._graph_key
    with ops.Graph().as_default():
      ops.get_default_graph()._graph_key = outside_graph_key
      if outside_executed_eagerly:
        with context.eager_mode():
          result = f(self, **kwargs)
      else:
        result = f(self, **kwargs)
    # Make an effort to clear caches, which would otherwise look like leaked
    # Tensors.
    context.context()._clear_caches()  # pylint: disable=protected-access
    gc.collect()
    tensors_after = [
        obj for obj in gc.get_objects()
        if _is_tensorflow_object(obj) and id(obj) not in tensors_before
    ]
    if tensors_after:
      raise AssertionError(("%d Tensors not deallocated after test: %s" % (
          len(tensors_after),
          str(tensors_after),
      )))
    return result

  return tf_decorator.make_decorator(f, decorator)


def _find_reference_cycle(objects: Sequence[Any], idx: int) -> bool:

  def get_ignore_reason(obj: Any, denylist: Collection[Any]) -> Optional[str]:
    """Tests whether an object should be omitted from the dependency graph."""
    if len(denylist) > 100:
      return "<depth limit>"
    if tf_inspect.isframe(obj):
      if "test_util.py" in tf_inspect.getframeinfo(obj)[0]:
        return "<test code>"
    for b in denylist:
      if b is obj:
        return "<test code>"
    if obj is denylist:
      return "<test code>"
    return None

  # Note: this function is meant to help with diagnostics. Its output is purely
  # a human-readable representation, so you may freely modify it to suit your
  # needs.
  def describe(
      obj: Any, denylist: Collection[Any], leaves_only: bool = False,
  ) -> str:
    """Returns a custom human-readable summary of obj.

    Args:
      obj: the value to describe.
      denylist: same as denylist in get_ignore_reason.
      leaves_only: boolean flag used when calling describe recursively. Useful
        for summarizing collections.
    """
    if get_ignore_reason(obj, denylist):
      return "{}{}".format(get_ignore_reason(obj, denylist), type(obj))
    if tf_inspect.isframe(obj):
      return "frame: {}".format(tf_inspect.getframeinfo(obj))
    elif tf_inspect.ismodule(obj):
      return "module: {}".format(obj.__name__)
    else:
      if leaves_only:
        return "{}, {}".format(type(obj), id(obj))
      elif isinstance(obj, list):
        return "list({}): {}".format(
            id(obj), [describe(e, denylist, leaves_only=True) for e in obj])
      elif isinstance(obj, tuple):
        return "tuple({}): {}".format(
            id(obj), [describe(e, denylist, leaves_only=True) for e in obj])
      elif isinstance(obj, dict):
        return "dict({}): {} keys".format(id(obj), len(obj.keys()))
      elif tf_inspect.isfunction(obj):
        return "function({}) {}; globals ID: {}".format(
            id(obj), obj.__name__, id(obj.__globals__))
      else:
        return "{}, {}".format(type(obj), id(obj))

  def build_ref_graph(
      obj: Any,
      graph: dict[int, list[int]],
      reprs: dict[int, str],
      denylist: tuple[Any, ...],
  ) -> None:
    """Builds a reference graph as <referrer> -> <list of referents>.

    Args:
      obj: The object to start from. The graph will be built by recursively
        adding its referrers.
      graph: Dict holding the graph to be built. To avoid creating extra
        references, the graph holds object IDs rather than actual objects.
      reprs: Auxiliary structure that maps object IDs to their human-readable
        description.
      denylist: List of objects to ignore.
    """
    referrers = gc.get_referrers(obj)
    denylist = denylist + (referrers,)

    obj_id = id(obj)
    for r in referrers:
      if get_ignore_reason(r, denylist) is None:
        r_id = id(r)
        if r_id not in graph:
          graph[r_id] = []
        if obj_id not in graph[r_id]:
          graph[r_id].append(obj_id)
          build_ref_graph(r, graph, reprs, denylist)
          reprs[r_id] = describe(r, denylist)

  def find_cycle(
      el: int,
      graph: dict[int, list[int]],
      reprs: dict[int, str],
      path: tuple[int, ...],
  ) -> Optional[bool]:
    """Finds and prints a single cycle in the dependency graph."""
    if el not in graph:
      return
    for r in graph[el]:
      if r in path:
        logging.error("Reference cycle sample:")
        for p in path + (r,):
          logging.error(reprs.get(p, "unknown object " + str(p)))
        return True
      else:
        if find_cycle(r, graph, reprs, path + (r,)):
          return True
    return False

  obj = objects[idx]
  graph: dict[int, list[int]] = {}  # referrer ID -> object ID
  reprs: dict[int, str] = {}  # object ID -> description
  build_ref_graph(obj, graph, reprs, (objects, graph, reprs, get_ignore_reason,
                                      describe, build_ref_graph, find_cycle))
  for k in graph:
    if find_cycle(k, graph, reprs, ()):
      return True
  return False


def assert_no_garbage_created(f: _F) -> _F:
  """Test method decorator to assert that no garbage has been created.

  Note that this decorator sets DEBUG_SAVEALL, which in some Python interpreters
  cannot be un-set (i.e. will disable garbage collection for any other unit
  tests in the same file/shard).

  Args:
    f: The function to decorate.

  Returns:
    The decorated function.
  """

  # FIXME(power) -- Update documentation, we no longer care if garbage is
  # created, we only want to verify we don't have memory leaks.
  def decorator(self: "TensorFlowTestCase", **kwargs):
    """Sets DEBUG_SAVEALL, runs the test, and checks for new garbage."""
    gc.disable()
    previous_debug_flags = gc.get_debug()
    gc.set_debug(gc.DEBUG_UNCOLLECTABLE)
    gc.collect()
    previous_garbage = len(gc.garbage)
    result = f(self, **kwargs)
    gc.collect()
    new_garbage = len(gc.garbage)
    if new_garbage > previous_garbage:

      for i, obj in enumerate(gc.garbage[previous_garbage:]):
        # Known false positive for ast.fix_missing_locations.
        if getattr(obj, "__module__", "") == "ast":
          new_garbage -= 3

    if new_garbage > previous_garbage:
      logging.error(
          "The decorated test created work for Python's garbage collector, "
          "likely due to a reference cycle. New objects in cycle(s):")
      for i, obj in enumerate(gc.garbage[previous_garbage:]):
        try:
          logging.error("Object %d of %d", i,
                        len(gc.garbage) - previous_garbage)

          def _safe_object_str(obj) -> str:
            return "<%s %d>" % (obj.__class__.__name__, id(obj))

          logging.error("  Object type: %s", _safe_object_str(obj))
          logging.error(
              "  Referrer types: %s", ", ".join(
                  [_safe_object_str(ref) for ref in gc.get_referrers(obj)]))
          logging.error(
              "  Referent types: %s", ", ".join(
                  [_safe_object_str(ref) for ref in gc.get_referents(obj)]))
          logging.error("  Object attribute names: %s", dir(obj))
          logging.error("  Object __str__:")
          logging.error(obj)
          logging.error("  Object __repr__:")
          logging.error(repr(obj))
        except Exception:  # pylint: disable=broad-except
          logging.error("(Exception while printing object)")

    # When garbage is created, this call can help identify reference cycles,
    # which are typically the cause of such garbage.
    if new_garbage > previous_garbage:
      for i in range(previous_garbage, new_garbage):
        if _find_reference_cycle(gc.garbage, i):
          break

    # This will fail if any garbage has been created, typically because of a
    # reference cycle.
    self.assertEqual(previous_garbage, new_garbage)
    # TODO(allenl): Figure out why this debug flag reset doesn't work. It would
    # be nice to be able to decorate arbitrary tests in a large test suite and
    # not hold on to every object in other tests.
    gc.set_debug(previous_debug_flags)
    gc.enable()
    return result

  return decorator


def _combine_named_parameters(**kwargs) -> list[OrderedDict[str, Any]]:
  """Generate combinations based on its keyword arguments.

  Two sets of returned combinations can be concatenated using +.  Their product
  can be computed using `times()`.

  Args:
    **kwargs: keyword arguments of form `option=[possibilities, ...]` or
      `option=the_only_possibility`.

  Returns:
    a list of dictionaries for each combination. Keys in the dictionaries are
    the keyword argument names.  Each key has one value - one of the
    corresponding keyword argument values.
  """
  sort_by_key = lambda k: k[0]
  combinations: list[list[tuple[str, Any]]] = []
  for key, values in sorted(kwargs.items(), key=sort_by_key):
    if not isinstance(values, list):
      values = [values]
    combinations.append([(key, value) for value in values])

  return [OrderedDict(result) for result in itertools.product(*combinations)]


def generate_combinations_with_testcase_name(
    **kwargs,
) -> list[OrderedDict[str, Any]]:
  """Generate combinations based on its keyword arguments using combine().

  This function calls combine() and appends a testcase name to the list of
  dictionaries returned. The 'testcase_name' key is a required for named
  parameterized tests.

  Args:
    **kwargs: keyword arguments of form `option=[possibilities, ...]` or
      `option=the_only_possibility`.

  Returns:
    a list of dictionaries for each combination. Keys in the dictionaries are
    the keyword argument names.  Each key has one value - one of the
    corresponding keyword argument values.
  """
  combinations = _combine_named_parameters(**kwargs)
  named_combinations: list[OrderedDict[str, Any]] = []
  for combination in combinations:
    assert isinstance(combination, OrderedDict)
    name = "".join([
        "_{}_{}".format("".join(filter(str.isalnum, key)),
                        "".join(filter(str.isalnum, str(value))))
        for key, value in combination.items()
    ])
    named_combinations.append(
        OrderedDict(
            list(combination.items()) +
            [("testcase_name", "_test{}".format(name))]))

  return named_combinations


def run_all_in_graph_and_eager_modes(cls: _TC) -> _TC:
  """Execute all test methods in the given class with and without eager."""
  base_decorator = run_in_graph_and_eager_modes
  for name in dir(cls):
    if (not name.startswith(unittest.TestLoader.testMethodPrefix) or
        name.startswith("testSkipEager") or
        name.startswith("test_skip_eager") or
        name == "test_session" or
        name == "test_scope"):
      continue
    value = getattr(cls, name, None)
    if callable(value):
      setattr(cls, name, base_decorator(value))
  return cls


def run_class_in_v1_v2(cls: _TC) -> _TC:
  """Execute all test methods in a given class in v1 and v2 modes."""
  base_decorator = run_in_v1_v2
  for name in dir(cls):
    if (not name.startswith(unittest.TestLoader.testMethodPrefix) or
        name.startswith("testSkipEager") or
        name.startswith("test_skip_eager") or
        name == "test_session" or
        name == "test_scope"):
      continue

    attr = getattr(cls, name, None)
    if not callable(attr):
      continue

    setattr(cls, name, base_decorator(attr))
  return cls


def enable_nested_function_shape_inference(fn: _F) -> _F:
  """Decorator for enabling nested_function_shape_inference on a test.

  This function returns a decorator intended to be applied to test methods in
  a `tf.test.TestCase` class. Doing so will set nested_function_shape_inference,
  reset the context, execute the test, then reset the context to the state
  it was in prior to this test.

  Example:

  class MyTest(test.TestCase):

    @enable_nested_function_shape_inference
    def testFoo(self):
      ...

  Args:
    fn: the function to be wrapped.

  Returns:
    The wrapped function.
  """

  def wrapper(*args, **kwargs):
    # If `nested_function_shape_inference` is already enabled do nothing.
    if flags.config().enable_nested_function_shape_inference.value():
      return fn(*args, **kwargs)

    flags.config().enable_nested_function_shape_inference.reset(True)
    try:
      return fn(*args, **kwargs)
    finally:
      flags.config().enable_nested_function_shape_inference.reset(False)

  return wrapper


def enable_quantized_dtypes_training(fn: _F) -> _F:
  """Decorator for enabling quantized_dtypes_training on a test.

  This function returns a decorator intended to be applied to test methods in
  a `tf.test.TestCase` class. Doing so will set quantized_dtypes_training,
  reset the context, execute the test, then reset the context to the state
  it was in prior to this test.

  Example:

  class MyTest(test.TestCase):

    @enable_quantized_dtypes_training
    def testFoo(self):
      ...

  Args:
    fn: the function to be wrapped.

  Returns:
    The wrapped function.
  """

  def wrapper(*args, **kwargs):
    # If `enable_quantized_dtypes_training` is already enabled do nothing.
    if flags.config().enable_quantized_dtypes_training.value():
      return fn(*args, **kwargs)

    flags.config().enable_quantized_dtypes_training.reset(True)
    try:
      return fn(*args, **kwargs)
    finally:
      flags.config().enable_quantized_dtypes_training.reset(False)

  return wrapper


def enable_eager_op_as_function(fn: _F) -> _F:
  """Returns the same fn. This will be removed once all usages are removed.

  Args:
    fn: the function to be wrapped.

  Returns:
    The wrapped function.
  """

  def wrapper(*args, **kwargs):
    return fn(*args, **kwargs)

  return wrapper


@overload
def with_eager_op_as_function(
    cls: type[_T],
    only_as_function: bool = False,
) -> type[_T]:
  ...


@overload
def with_eager_op_as_function(
    cls: None = None,
    only_as_function: bool = False,
) -> Callable[[type[_T]], type[_T]]:
  ...


@tf_export("test.with_eager_op_as_function")
def with_eager_op_as_function(
    cls: Optional[type[_T]] = None,
    only_as_function: bool = False,  # pylint: disable=unused-argument
) -> Union[Callable[[type[_T]], type[_T]], type[_T]]:
  """Returns the same class. This will be removed once all usages are removed.

  Args:
    cls: class to decorate.
    only_as_function: unused argument.

  Returns:
    cls
  """

  def decorator(cls: type[_T]) -> type[_T]:
    return cls

  if cls is not None:
    return decorator(cls)

  return decorator  # pytype: disable=bad-return-type


def enable_graph_building_optimization(fn: _F) -> _F:
  """Decorator for enabling graph_building_optimization on a test.

  This function returns a decorator intended to be applied to test methods in
  a `tf.test.TestCase` class. Doing so will enable graph_building_optimization,
  execute the test, then reset the feature flag to its default value.

  Example:

  class MyTest(test.TestCase):

    @enable_graph_building_optimization
    def testFoo(self):
      ...

  Args:
    fn: the function to be wrapped.

  Returns:
    The wrapped function.
  """

  def wrapper(*args, **kwargs):
    # If `graph_building_optimization` is already enabled do nothing.
    if flags.config().graph_building_optimization.value():
      return fn(*args, **kwargs)

    flags.config().graph_building_optimization.reset(True)
    try:
      return fn(*args, **kwargs)
    finally:
      flags.config().graph_building_optimization.reset(False)

  return wrapper


def add_graph_building_optimization_tests(cls: _TC) -> _TC:
  """Adds methods with graph_building_optimization enabled to the test suite.

  Example:

  @test_util.add_graph_building_optimization_tests
  class FooTest(test.TestCase):

    def testBar(self):
      ...

  Generated class:
  class FooTest(test.TestCase):

    def testBar(self):
      ...

    def testBarWithGraphBuildingOptimization(self):
      // Enable graph_building_optimization
      testBar(self)
      // Disable graph_building_optimization

  Args:
    cls: class to decorate.

  Returns:
    cls with new test methods added.
  """

  if flags.config().graph_building_optimization.value():
    return cls

  for name, value in cls.__dict__.copy().items():
    if (callable(value) and
        (name.startswith(unittest.TestLoader.testMethodPrefix) or
         name.startswith("benchmark"))):
      setattr(cls, name + "WithGraphBuildingOptimization",
              enable_graph_building_optimization(value))
  return cls


def disable_eager_op_as_function(unused_msg: str) -> Callable[[_F], _F]:
  """Decorator for a function in a with_eager_op_as_function enabled test class.

  Blocks the function from being run with eager_op_as_function enabled.

  Args:
    unused_msg: Reason for disabling.

  Returns:
    The wrapped function with _disable_eager_op_as_function attr set to True.
  """
  return _disable_test(execute_func=False)


def set_xla_env_flag(flag: str = "") -> Callable[[_F], _F]:
  """Decorator for setting XLA_FLAGS prior to running a test.

  This function returns a decorator intended to be applied to test methods in
  a `tf.test.TestCase` class. Doing so will allow users to set any xla flags
  exposed via the XLA_FLAGS environment variable, execute the test, then reset
  the XLA_FLAGS to the state it was in prior to this test.

  Example:

  class MyTest(test.TestCase):

    @set_xla_env_flag(flag='--xla_gpu_enable_fast_min_max=false')
    def testFoo(self):
      ...

  Args:
    flag: The xla flag to be set in the XLA_FLAGS env variable.

  Returns:
    A decorator which sets the configured flag in XLA_FLAGS for the decorated
    function.
  """

  def decorator(f: _F) -> _F:

    @functools.wraps(f)
    def decorated(*args, **kwargs):
      original_xla_flags = os.environ.get("XLA_FLAGS")
      new_xla_flags = flag
      if original_xla_flags:
        new_xla_flags = new_xla_flags + " " + original_xla_flags
      os.environ["XLA_FLAGS"] = new_xla_flags
      try:
        return f(*args, **kwargs)
      finally:
        if original_xla_flags is None:
          del os.environ["XLA_FLAGS"]
        else:
          os.environ["XLA_FLAGS"] = original_xla_flags

    return decorated

  return decorator


def build_as_function_and_v1_graph(
    func: Callable[..., Any],
) -> Callable[..., None]:
  """Run a test case in v1 graph mode and inside tf.function in eager mode.

  WARNING: This decorator can only be used in test cases that statically checks
  generated graph. Attempting to evaluate graph or function results via.
  session.run() or self.evaluate() will fail.

  WARNING: This decorator can only be used for test cases that inherit from
  absl.testing.parameterized.TestCase.

  Args:
    func: Test case function to be decorated.

  Returns:
    Decorated test case function.
  """

  if tf_inspect.isclass(func):
    raise ValueError(
        "`run_in_graph_mode_and_function` only supports test methods.")

  @parameterized.named_parameters(("_v1_graph", "v1_graph"),
                                  ("_function", "function"))
  @functools.wraps(func)
  def decorated(
      self: "TensorFlowTestCase",
      run_mode: str,
      *args,
      **kwargs,
  ) -> None:
    if run_mode == "v1_graph":
      with ops.Graph().as_default():
        func(self, *args, **kwargs)
    elif run_mode == "function":

      @def_function.function
      def function_in_eager():
        func(self, *args, **kwargs)

      # Create a new graph for the eagerly executed version of this test for
      # better isolation.
      graph_for_eager_test = ops.Graph()
      with graph_for_eager_test.as_default(), context.eager_mode():
        function_in_eager()
      ops.dismantle_graph(graph_for_eager_test)
    else:
      raise ValueError("Unknown run mode %s" % run_mode)

  return decorated


def run_in_async_and_sync_mode(f: _F) -> _F:
  """Execute the test in async mode and sync mode."""

  @parameterized.named_parameters([("Async", True), ("", False)])
  @functools.wraps(f)
  def decorator(self: "TensorFlowTestCase", async_mode: bool, *args, **kwargs):
    if async_mode:
      with context.execution_mode(context.ASYNC):
        f(self, *args, **kwargs)
    else:
      with context.execution_mode(context.SYNC):
        f(self, *args, **kwargs)
  return decorator


@overload
def run_in_graph_and_eager_modes(
    func: Callable[..., Any],
    config: Optional[config_pb2.ConfigProto] = None,
    use_gpu: bool = True,
    assert_no_eager_garbage: bool = False,
) -> Callable[..., None]:
  ...


@overload
def run_in_graph_and_eager_modes(
    func: None = None,
    config: Optional[config_pb2.ConfigProto] = None,
    use_gpu: bool = True,
    assert_no_eager_garbage: bool = False,
) -> Callable[[Callable[..., Any]], Callable[..., None]]:
  ...


def run_in_graph_and_eager_modes(
    func: Optional[Callable[..., Any]] = None,
    config: Optional[config_pb2.ConfigProto] = None,
    use_gpu: bool = True,
    assert_no_eager_garbage: bool = False,
) -> Union[
    Callable[[Callable[..., Any]], Callable[..., None]],
    Callable[..., None],
]:
  """Execute the decorated test with and without enabling eager execution.

  This function returns a decorator intended to be applied to test methods in
  a `tf.test.TestCase` class. Doing so will cause the contents of the test
  method to be executed twice - once normally, and once with eager execution
  enabled. This allows unittests to confirm the equivalence between eager
  and graph execution (see `tf.compat.v1.enable_eager_execution`).

  For example, consider the following unittest:

  ```python
  class MyTests(tf.test.TestCase):

    @run_in_graph_and_eager_modes
    def test_foo(self):
      x = tf.constant([1, 2])
      y = tf.constant([3, 4])
      z = tf.add(x, y)
      self.assertAllEqual([4, 6], self.evaluate(z))

  if __name__ == "__main__":
    tf.test.main()
  ```

  This test validates that `tf.add()` has the same behavior when computed with
  eager execution enabled as it does when constructing a TensorFlow graph and
  executing the `z` tensor in a session.

  `deprecated_graph_mode_only`, `run_v1_only`, `run_v2_only`, and
  `run_in_graph_and_eager_modes` are available decorators for different
  v1/v2/eager/graph combinations.


  Args:
    func: function to be annotated. If `func` is None, this method returns a
      decorator the can be applied to a function. If `func` is not None this
      returns the decorator applied to `func`.
    config: An optional config_pb2.ConfigProto to use to configure the session
      when executing graphs.
    use_gpu: If True, attempt to run as many operations as possible on GPU.
    assert_no_eager_garbage: If True, sets DEBUG_SAVEALL on the garbage
      collector and asserts that no extra garbage has been created when running
      the test with eager execution enabled. This will fail if there are
      reference cycles (e.g. a = []; a.append(a)). Off by default because some
      tests may create garbage for legitimate reasons (e.g. they define a class
      which inherits from `object`), and because DEBUG_SAVEALL is sticky in some
      Python interpreters (meaning that tests which rely on objects being
      collected elsewhere in the unit test file will not work). Additionally,
      checks that nothing still has a reference to Tensors that the test
      allocated.

  Returns:
    Returns a decorator that will run the decorated test method twice:
    once by constructing and executing a graph in a session and once with
    eager execution enabled.
  """

  def decorator(f: Callable[..., Any]) -> Callable[..., None]:
    if tf_inspect.isclass(f):
      raise ValueError(
          "`run_in_graph_and_eager_modes` only supports test methods. "
          "Did you mean to use `run_all_in_graph_and_eager_modes`?")

    def decorated(self: "TensorFlowTestCase", *args, **kwargs) -> None:
      logging.info("Running %s in GRAPH mode.", f.__name__)
      try:
        with context.graph_mode(), self.subTest("graph_mode"):
          # XLATestCase uses `session`, which also doesn't take any args,
          # instead of `test_session`
          for class_ in self.__class__.mro():
            if class_.__name__ == "XLATestCase":
              session_func = self.session
              session_kwargs = {}
              break
          else:
            session_func = self.test_session
            session_kwargs = dict(use_gpu=use_gpu, config=config)
          with session_func(**session_kwargs):
            f(self, *args, **kwargs)
      except unittest.case.SkipTest:
        pass

      def run_eagerly(self: "TensorFlowTestCase", **kwargs) -> None:
        logging.info("Running %s in EAGER mode.", f.__name__)
        if not use_gpu:
          with ops.device("/device:CPU:0"):
            f(self, *args, **kwargs)
        else:
          f(self, *args, **kwargs)

      if assert_no_eager_garbage:
        ops.reset_default_graph()
        run_eagerly = assert_no_new_tensors(
            assert_no_garbage_created(run_eagerly))

      # This decorator runs the wrapped test twice.
      # Reset the test environment between runs.
      self.tearDown()
      self._tempdir = None
      # Create a new graph for the eagerly executed version of this test for
      # better isolation.
      graph_for_eager_test = ops.Graph()
      with (
          graph_for_eager_test.as_default(),
          context.eager_mode(),
          self.subTest("eager_mode"),
      ):
        self.setUp()
        run_eagerly(self, **kwargs)
      ops.dismantle_graph(graph_for_eager_test)

    return tf_decorator.make_decorator(f, decorated)

  if func is not None:
    return decorator(func)

  return decorator


def run_in_v1_v2(
    device_to_use: Optional[str] = None,
    assert_no_eager_garbage: bool = False,
) -> Callable[[Callable[..., Any]], Callable[..., None]]:
  """Execute the decorated test in v1 and v2 modes.

  The overall execution is similar to that of `run_in_graph_and_eager_mode`.

  Args:
    device_to_use: A string in the following format: "/device:CPU:0".
    assert_no_eager_garbage: If True, sets DEBUG_SAVEALL on the garbage
      collector and asserts that no extra garbage has been created when running
      the test with eager execution enabled. This will fail if there are
      reference cycles (e.g. a = []; a.append(a)). Off by default because some
      tests may create garbage for legitimate reasons (e.g. they define a class
      which inherits from `object`), and because DEBUG_SAVEALL is sticky in some
      Python interpreters (meaning that tests which rely on objects being
      collected elsewhere in the unit test file will not work). Additionally,
      checks that nothing still has a reference to Tensors that the test
      allocated.

  Returns:
    A decorator that runs a given test in v1 and v2 modes.
  """

  def decorator(f: Callable[..., Any]) -> Callable[..., None]:
    decorator_tag = "wrapped_with_v1_v2_decorator"
    if hasattr(f, decorator_tag):
      # Already decorated with this very same decorator
      return f

    def decorated(self: "TensorFlowTestCase", *args, **kwargs) -> None:
      logging.info("Running %s in V1 mode.", f.__name__)
      try:
        with self.subTest("V1_mode"):
          v2_compat.disable_v2_behavior()
          f(self, *args, **kwargs)
      except unittest.case.SkipTest:
        pass

      def run_v2(self: "TensorFlowTestCase", **kwargs) -> None:
        logging.info("Running %s in V2 mode.", f.__name__)
        if device_to_use:
          with ops.device(device_to_use):
            f(self, *args, **kwargs)
        else:
          f(self, *args, **kwargs)

      if assert_no_eager_garbage:
        ops.reset_default_graph()
        run_v2 = assert_no_new_tensors(
            assert_no_garbage_created(run_v2))

      # This decorator runs the wrapped test twice.
      # Reset the test environment between runs.
      self.tearDown()
      self._tempdir = None  # pylint:disable=protected-access

      ops.reset_default_graph()
      v2_compat.enable_v2_behavior()
      with self.subTest("V2_mode"):
        self.setUp()
        run_v2(self, **kwargs)

    tf_decorated = tf_decorator.make_decorator(f, decorated)
    tf_decorated.__dict__[decorator_tag] = True
    return tf_decorated

  return decorator


def py_func_if_in_function(f: _F) -> _F:

  def decorated(*args, **kwds):
    if not ops.inside_function():
      return f(*args, **kwds)

    tensor_args: list[Union[tensor_lib.Tensor, variables.Variable]] = []
    tensor_indices: list[int] = []
    for i, arg in enumerate(args):
      if isinstance(arg, (tensor_lib.Tensor, variables.Variable)):
        tensor_args.append(arg)
        tensor_indices.append(i)

    def inner_f(*inner_tensor_args):
      my_args = list(args)
      for i, n in zip(tensor_indices, inner_tensor_args):
        my_args[i] = n
      return f(*my_args, **kwds)

    return script_ops.py_func(inner_f, tensor_args, [])

  return tf_decorator.make_decorator(f, decorated)


def also_run_as_tf_function(f: Callable[..., Any]) -> Callable[..., None]:
  """Runs the decorated test twice--once as is, once inside a tf.function.

  This allows you to run a test both in eager execution and inside a
  tf.function, exercising the two execution modes supported in tf 2.0. The test
  assertions are automatically done inside tf.py_funcs, and tf.function ensures
  that they run in the proper order and with the proper side effects.

  Currently variable creation is not supported in tests annotated with this
  decorator since it's tricky to ensure the variable doesn't get repeatedly
  created when retracing the tf.function.

  Args:
    f: the test method to be decorated

  Returns:
    The decorated test method, which will run both in eager and inside a
    tf.function.
  """

  def decorated(*args, **kwds) -> None:

    def bound_f() -> None:
      f(*args, **kwds)

    with context.eager_mode():
      # Running in eager mode
      bound_f()
      # Running as TF function
      # TODO(b/121143941): Remove the autograph override.
      def_function.function(bound_f, autograph=False)()

  return decorated


@overload
def deprecated_graph_mode_only(func: _F) -> _F:
  ...


@overload
def deprecated_graph_mode_only(func: _TC) -> Optional[_TC]:
  ...


def deprecated_graph_mode_only(func: Union[_TC, _F]) -> Union[_TC, _F]:
  """Execute the decorated test in graph mode.

  This is a decorator intended to be applied to tests that are not compatible
  with eager mode. When this decorator is applied, the test body will be run in
  an environment where API calls construct graphs instead of executing eagerly.

  `deprecated_graph_mode_only`, `run_v1_only`, `run_v2_only`, and
  `run_in_graph_and_eager_modes` are available decorators for different
  v1/v2/eager/graph combinations.

  Args:
    func: function or class to be annotated.
      If `func` is a function this returns the decorator applied to `func`.
      If `func` is a unit test class this returns that class with the decorator
      applied to all test functions within that class.

  Returns:
    Returns a function or class that will run the decorated test(s)
    in graph mode.
  """

  if tf_inspect.isclass(func):
    setup = func.__dict__.get("setUp")
    if setup is not None:
      setattr(func, "setUp", deprecated_graph_mode_only(setup))

    for name, value in func.__dict__.copy().items():
      if (callable(value) and
          name.startswith(unittest.TestLoader.testMethodPrefix)):
        setattr(func, name, deprecated_graph_mode_only(value))

    return func

  def decorated(*args, **kwargs):
    if context.executing_eagerly():
      with context.graph_mode():
        return func(*args, **kwargs)
    else:
      return func(*args, **kwargs)

  return tf_decorator.make_decorator(func, decorated)


run_deprecated_v1 = deprecated_graph_mode_only


def run_all_in_deprecated_graph_mode_only(cls: _TC) -> _TC:
  """Execute all tests in a class in graph mode."""
  base_decorator = deprecated_graph_mode_only
  for name in dir(cls):
    if (not name.startswith(unittest.TestLoader.testMethodPrefix) or
        name == "test_session"):
      continue
    value = getattr(cls, name, None)
    if callable(value):
      setattr(cls, name, base_decorator(value))
  return cls


def _run_vn_only(func=None, v2=True, reason=None):
  """Execute the decorated test only if running in the specified mode.

  This function is intended to be applied to tests that exercise functionality
   that belongs to either only v2, or v1.
   If the test is run in the mode opposite of the specified one, it will simply
   be skipped.

   It shouldn't be used directly, instead, use the `run_v1_only` or
   `run_v2_only` wrappers that call it.

  `deprecated_graph_mode_only`, `run_v1_only`, `run_v2_only`, and
  `run_in_graph_and_eager_modes` are available decorators for different
  v1/v2/eager/graph combinations.

  Args:
    func: function to be annotated. If `func` is None, this method returns a
      decorator the can be applied to a function. If `func` is not None this
      returns the decorator applied to `func`.
    v2: a boolean value indicating whether the test should be skipped in v2,
      or v1.
    reason: string giving a reason for limiting the test to a particular mode.

  Returns:
    A decorator that will skip the test method in the specified version.
  """
  if not reason:
    reason = f"Test is only compatible with {'v2 ' if v2 else 'v1'}"

  def decorator(f):
    if tf_inspect.isclass(f):
      # To skip an entire test suite class, we only decorate the setUp method
      # to skip all tests. There are cases when setUp is not defined (not
      # overridden in subclasses of TestCase, so not available in f.__dict__
      # below). For those cases, we walk the method resolution order list and
      # pick the first setUp method we find (usually this should be the one in
      # the parent class since that's the TestCase class).
      for cls in type.mro(f):
        setup = cls.__dict__.get("setUp")
        if setup is not None:
          setattr(f, "setUp", decorator(setup))
          break

      return f
    else:
      # If f is just a function, just create a decorator for it and return it
      def decorated(self, *args, **kwargs):
        tf2_enabled = tf2.enabled()
        # Skip if TF is in v2 mode, but the test is expected to only be run
        # in v1, and vice versa
        if (tf2_enabled and not v2) or (not tf2_enabled and v2):
          self.skipTest(reason)

        return f(self, *args, **kwargs)

      return tf_decorator.make_decorator(f, decorated)

  if func is not None:
    return decorator(func)

  return decorator


def run_v1_only(reason=None, func=None):
  """Only execute the test if Tensorflow is in v1 mode."""
  return _run_vn_only(func=func, v2=False, reason=reason)


def run_v2_only(func=None, reason=None):
  """Only execute the test if Tensorflow is in v2 mode."""
  return _run_vn_only(func=func, v2=True, reason=reason)


def run_gpu_only(func: _F) -> _F:
  """Execute the decorated test only if a GPU is available.

  This function is intended to be applied to tests that require the presence
  of a GPU. If a GPU is absent, it will simply be skipped.

  Args:
    func: function to be annotated.

  Returns:
    Returns a function that will conditionally skip the decorated test method.
  """

  if tf_inspect.isclass(func):
    raise ValueError("`run_gpu_only` only supports test methods.")

  def decorated(self: "TensorFlowTestCase", *args, **kwargs):
    if not is_gpu_available():
      self.skipTest("Test requires GPU")

    return func(self, *args, **kwargs)

  return decorated


def run_cuda_only(func: _F) -> _F:
  """Execute the decorated test only if a GPU is available.

  This function is intended to be applied to tests that require the presence
  of a CUDA GPU. If a CUDA GPU is absent, it will simply be skipped.

  Args:
    func: function to be annotated.

  Returns:
    Returns a function that will conditionally skip the decorated test method.
  """

  if tf_inspect.isclass(func):
    raise ValueError("`run_cuda_only` only supports test methods.")

  def decorated(self: "TensorFlowTestCase", *args, **kwargs):
    if not is_gpu_available(cuda_only=True):
      self.skipTest("Test requires CUDA GPU")

    return func(self, *args, **kwargs)

  return decorated


def run_gpu_or_tpu(func: _F) -> _F:
  """Execute the decorated test only if a physical GPU or TPU is available.

  This function is intended to be applied to tests that require the presence
  of a physical GPU or TPU. It complies with the following rules:
  - If a GPU is available, the test will run on the GPU.
  - If a GPU is absent and a TPU is available, the test will run on the TPU.
  - If both GPU and TPU are absent, the test will be skipped.

  Args:
    func: function to be annotated.

  Returns:
    Returns a function that will conditionally skip the decorated test method.
  """

  if tf_inspect.isclass(func):
    raise ValueError("`run_gpu_or_tpu` only supports test methods.")

  def decorated(self: "TensorFlowTestCase", *args, **kwargs):
    if config.list_physical_devices("GPU"):
      return func(self, "GPU", *args, **kwargs)

    if config.list_physical_devices("TPU"):
      return func(self, "TPU", *args, **kwargs)

    self.skipTest("Test requires GPU or TPU")

  return decorated


def with_forward_compatibility_horizons(
    *horizons: Optional[tuple[int, int, int]]
) -> Callable[[Callable[..., Any]], Callable[..., None]]:
  """Executes the decorated test with the specified forward-compat horizons.

  Args:
    *horizons: A list of (year, month, day) tuples.  If the list includes
      `None`, then the test will also be run with no forward-compatibility
      horizon set.

  Returns:
    A decorator that will execute the test with the specified horizons.
  """
  if not horizons:
    raise ValueError("Expected at least one horizon.")
  for horizon in horizons:
    if not ((horizon is None) or
            (len(horizon) == 3 and all(isinstance(x, int) for x in horizon))):
      raise ValueError("Bad horizon value: %r" % horizon)

  def decorator(f: Callable[..., Any]) -> Callable[..., None]:
    if tf_inspect.isclass(f):
      raise ValueError("`with_forward_compatibility_horizons` only "
                       "supports test methods.")
    def decorated(*args, **kwargs):
      for horizon in horizons:
        if horizon is None:
          f(*args, **kwargs)
        else:
          (year, month, day) = horizon
          with forward_compatibility_horizon(year, month, day):
            f(*args, **kwargs)
    return tf_decorator.make_decorator(f, decorated)

  return decorator


@deprecation.deprecated(None,
                        "Use `tf.config.list_physical_devices('GPU')` instead.")
@tf_export("test.is_gpu_available")
def is_gpu_available(
    cuda_only: bool = False,
    min_cuda_compute_capability: Optional[tuple[int, int]] = None,
) -> bool:
  """Returns whether TensorFlow can access a GPU.

  Warning: if a non-GPU version of the package is installed, the function would
  also return False. Use `tf.test.is_built_with_cuda` to validate if TensorFlow
  was build with CUDA support.

  For example,
  >>> gpu_available = tf.test.is_gpu_available()
  >>> is_cuda_gpu_available = tf.test.is_gpu_available(cuda_only=True)
  >>> is_cuda_gpu_min_3 = tf.test.is_gpu_available(True, (3,0))

  Args:
    cuda_only: limit the search to CUDA GPUs.
    min_cuda_compute_capability: a (major,minor) pair that indicates the minimum
      CUDA compute capability required, or None if no requirement.

  Note that the keyword arg name "cuda_only" is misleading (since routine will
  return true when a GPU device is available irrespective of whether TF was
  built with CUDA support or ROCm support. However no changes here because

  ++ Changing the name "cuda_only" to something more generic would break
     backward compatibility

  ++ Adding an equivalent "rocm_only" would require the implementation check
     the build type. This in turn would require doing the same for CUDA and thus
     potentially break backward compatibility

  ++ Adding a new "cuda_or_rocm_only" would not break backward compatibility,
     but would require most (if not all) callers to update the call to use
     "cuda_or_rocm_only" instead of "cuda_only"

  Returns:
    True if a GPU device of the requested kind is available.
  """

  # This was needed earlier when we had support for SYCL in TensorFlow.
  del cuda_only

  try:
    for local_device in device_lib.list_local_devices():
      if local_device.device_type == "GPU":
        gpu_info = gpu_util.compute_capability_from_device_desc(local_device)
        cc = gpu_info.compute_capability or (0, 0)
        if not min_cuda_compute_capability or cc >= min_cuda_compute_capability:
          return True
    return False
  except errors_impl.NotFoundError as e:
    if not all(x in str(e) for x in ["CUDA", "not find"]):
      raise e
    else:
      logging.error(str(e))
      return False


@contextlib.contextmanager
def device(use_gpu: bool) -> Iterator[None]:
  """Uses gpu when requested and available."""
  if use_gpu and is_gpu_available():
    dev = "/device:GPU:0"
  else:
    dev = "/device:CPU:0"
  with ops.device(dev):
    yield


@contextlib.contextmanager
def use_gpu() -> Iterator[None]:
  """Uses gpu when requested and available."""
  with device(use_gpu=True):
    yield


@contextlib.contextmanager
def force_gpu() -> Iterator[None]:
  """Force the gpu to be used."""
  with ops.device("/device:GPU:0"):
    yield


@contextlib.contextmanager
def force_cpu() -> Iterator[None]:
  """Force the cpu to be used."""
  with ops.device("/device:CPU:0"):
    yield


@contextlib.contextmanager
def deterministic_ops() -> Iterator[None]:
  """Enables deterministic ops."""
  try:
    config.enable_op_determinism()
    yield
  finally:
    config.disable_op_determinism()


class CapturedWrites:
  """A utility class to load the captured writes made to a stream."""

  def __init__(self, capture_location: str):
    self.capture_location = capture_location

  def contents(self) -> str:
    """Get the captured writes as a single string."""
    with open(self.capture_location) as tmp_file:
      output_data = "".join(tmp_file.readlines())
    return output_data


class FakeEagerSession:
  """Fake session so tests that conditionally use placeholders can use eager.

  There are a number of tests that conditionally use placeholders for shape
  inference. The pattern is demonstrated here:

  ```python
  with self.cached_session() as sess:
    if static_shape:
      y = math_ops.matmul(x, ...)
      feed_dict = {}
    else:
      x_ph = array_ops.placeholder(...)
      y = math_ops.matmul(x_ph, ...)
      feed_dict = {x_ph: x}
    val = sess.run(y, feed_dict=feed_dict)
  ```

  Since the feed_dict is empty when not using placeholders we should be able to
  call self.evaluate(), however this requires rewriting the test case.
  This class should be considered a stop-gap solution to get tests running with
  eager with minimal changes to the actual test.
  """

  def __init__(self, test_case):
    self._test_case = test_case

  def run(self, fetches, *args, **kwargs):
    """Evaluate `fetches`.

    Fail if additional args are specified.

    Args:
      fetches: A Tensor or a nested list/tuple of Tensors.
      *args: Positional arguments
      **kwargs: Keyword arguments

    Raises:
      RuntimeError: If args or kwargs are specified.

    Returns:
      Tensors as numpy values.
    """
    feed_dict = kwargs.pop("feed_dict", {})
    if feed_dict:
      raise RuntimeError(
          "feed_dict is not supported when eager execution is enabled "
          "(in this case, sess.run(t) is shorthand for t.numpy()")

    if args or kwargs:
      raise RuntimeError(
          "Optional args are not supported when eager execution is enabled "
          "(in this case, sess.run(t) is shorthand for t.numpy()")

    return self._test_case.evaluate(fetches)


class ErrorLoggingSession(s.Session):
  """Wrapper around a Session that logs errors in run()."""

  def run(self, *args, **kwargs):
    try:
      return super().run(*args, **kwargs)
    except Exception as e:  # pylint: disable=broad-except
      # Note: disable the logging for OutOfRangeError, which makes the output
      # of tf.data tests hard to read, because OutOfRangeError is used as the
      # signal completion
      if not isinstance(e, errors.OutOfRangeError):
        logging.error(str(e))
      raise


def disable_cudnn_autotune(func: _F) -> _F:
  """Disable autotuning during the call to this function.

  Some tests want to base assertions on a graph being isomorphic with a copy.
  To ensure this, this decorator disables autotuning.

  Args:
    func: Function to run with CuDNN autotuning turned off.

  Returns:
    Decorated function.
  """

  def decorated(*args, **kwargs):
    original_tf_cudnn_use_autotune = os.environ.get("TF_CUDNN_USE_AUTOTUNE")
    os.environ["TF_CUDNN_USE_AUTOTUNE"] = "false"
    original_xla_flags = os.environ.get("XLA_FLAGS")
    new_xla_flags = "--xla_gpu_autotune_level=0"
    if original_xla_flags:
      new_xla_flags = original_xla_flags + " " + new_xla_flags
    os.environ["XLA_FLAGS"] = new_xla_flags

    result = func(*args, **kwargs)

    if (original_tf_cudnn_use_autotune is None):
      del os.environ["TF_CUDNN_USE_AUTOTUNE"]
    else:
      os.environ["TF_CUDNN_USE_AUTOTUNE"] = original_tf_cudnn_use_autotune
    if (original_xla_flags is None):
      del os.environ["XLA_FLAGS"]
    else:
      os.environ["XLA_FLAGS"] = original_xla_flags

    return result

  return tf_decorator.make_decorator(func, decorated)


# The description is just for documentation purposes.
def enable_tf_xla_constant_folding(description: str) -> Callable[[_F], _F]:

  if not isinstance(description, str):
    raise ValueError("'description' should be string, got {}".format(
        type(description)))

  def enable_tf_xla_constant_folding_impl(func: _F) -> _F:
    """Enable constant folding during the call to this function.

    Some tests fail without constant folding.

    Args:
      func: Function to run with constant folding turned on.

    Returns:
      Decorated function.
    """

    def decorated(*args, **kwargs):
      original_var = pywrap_tf_session.TF_GetXlaConstantFoldingDisabled()
      pywrap_tf_session.TF_SetXlaConstantFoldingDisabled(False)
      result = func(*args, **kwargs)
      pywrap_tf_session.TF_SetXlaConstantFoldingDisabled(original_var)
      return result

    return tf_decorator.make_decorator(func, decorated)

  return enable_tf_xla_constant_folding_impl


# Updates test function by selectively disabling it.
def _disable_test(execute_func: bool) -> Callable[[_F], _F]:

  def disable_test_impl(func: _F) -> _F:

    def decorated(*args, **kwargs):
      if execute_func:
        return func(*args, **kwargs)

    return tf_decorator.make_decorator(func, decorated)

  return disable_test_impl


# The description is just for documentation purposes.
def disable_xla(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """Execute the test method only if xla is not enabled."""
  execute_func = not is_xla_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def disable_mlir_bridge(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """Execute the test method only if MLIR bridge is not enabled."""
  execute_func = not is_mlir_bridge_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def disable_asan(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """Execute the test method only if ASAN is not enabled."""
  execute_func = not is_asan_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def disable_msan(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """Execute the test method only if MSAN is not enabled."""
  execute_func = not is_msan_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def disable_tsan(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """Execute the test method only if TSAN is not enabled."""
  execute_func = not is_tsan_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def disable_ubsan(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """Execute the test method only if UBSAN is not enabled."""
  execute_func = not is_ubsan_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def disable_tfrt(
    unused_description: str,  # pylint: disable=unused-argument
) -> Callable[[Union[_TC, _F]], Union[_TC, _F, None]]:

  def disable_tfrt_impl(cls_or_func: Union[_TC, _F]) -> Union[_TC, _F, None]:
    """Execute the test only if tfrt is not enabled."""

    if tf_inspect.isclass(cls_or_func):
      if tfrt_utils.enabled():
        return None
      else:
        return cast(_TC, cls_or_func)
    else:
      func = cast(Callable[..., Any], cls_or_func)
      def decorated(*args, **kwargs):
        if tfrt_utils.enabled():
          return
        else:
          return func(*args, **kwargs)

      return tf_decorator.make_decorator(cls_or_func, decorated)

  return disable_tfrt_impl


def for_all_test_methods(
    decorator: Callable[..., Any], *args, **kwargs,
) -> Callable[[_TC], _TC]:
  """Generate class-level decorator from given method-level decorator.

  It is expected for the given decorator to take some arguments and return
  a method that is then called on the test method to produce a decorated
  method.

  Args:
    decorator: The decorator to apply.
    *args: Positional arguments
    **kwargs: Keyword arguments
  Returns: Function that will decorate a given classes test methods with the
    decorator.
  """

  def all_test_methods_impl(cls: _TC) -> _TC:
    """Apply decorator to all test methods in class."""
    for name in dir(cls):
      value = getattr(cls, name)
      if callable(value) and name.startswith(
          "test") and (name != "test_session"):
        setattr(cls, name, decorator(*args, **kwargs)(value))
    return cls

  return all_test_methods_impl


# The description is just for documentation purposes.
def no_xla_auto_jit(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument
  """This test is not intended to be run with XLA auto jit enabled."""
  execute_func = not is_xla_enabled()
  return _disable_test(execute_func)


# The description is just for documentation purposes.
def xla_allow_fallback(description: str) -> Callable[[_F], _F]:  # pylint: disable=unused-argument

  def xla_allow_fallback_impl(func: _F) -> _F:
    """Allow fallback to TF even though testing xla."""

    def decorated(*args, **kwargs):
      if is_xla_enabled():
        # Update the global XLABuildOpsPassFlags to enable lazy compilation,
        # which allows the compiler to fall back to TF classic. Remember the
        # old value so that we can reset it.
        old_value = pywrap_tf_session.TF_SetXlaEnableLazyCompilation(True)
        result = func(*args, **kwargs)
        pywrap_tf_session.TF_SetXlaEnableLazyCompilation(old_value)
        return result
      else:
        return func(*args, **kwargs)

    return tf_decorator.make_decorator(func, decorated)

  return xla_allow_fallback_impl


# The description is just for documentation purposes.
def run_without_tensor_float_32(
    description: str,  # pylint: disable=unused-argument
) -> Callable[[Callable[..., Any]], Callable[..., None]]:
  """Execute test with TensorFloat-32 disabled.

  While almost every real-world deep learning model runs fine with
  TensorFloat-32, many tests use assertAllClose or similar methods.
  TensorFloat-32 matmuls typically will cause such methods to fail with the
  default tolerances.

  Args:
    description: A description used for documentation purposes, describing why
      the test requires TensorFloat-32 to be disabled.

  Returns:
    Decorator which runs a test with TensorFloat-32 disabled.
  """

  def decorator(f: Callable[..., Any]) -> Callable[..., None]:

    @functools.wraps(f)
    def decorated(*args, **kwargs):
      allowed = config.tensor_float_32_execution_enabled()
      try:
        config.enable_tensor_float_32_execution(False)
        f(*args, **kwargs)
      finally:
        config.enable_tensor_float_32_execution(allowed)

    return tf_decorator.make_decorator(f, decorated)

  return decorator


# The description is just for documentation purposes.
def run_all_without_tensor_float_32(description: str) -> Callable[[_TC], _TC]:  # pylint: disable=unused-argument
  """Execute all tests in a class with TensorFloat-32 disabled."""
  return for_all_test_methods(run_without_tensor_float_32, description)


def matmul_without_tf32(a, b, *args, **kwargs):
  """Run matmul but cast float32 inputs to float64 if TensorFloat-32 is enabled.

  This effectively runs matmul without TensorFloat-32. It should only be used in
  tests when verifying some other op or functions works correctly, e.g. to test
  `tf.linalg.sqrtm` by matrix multiplying the output of the op by itself. In
  such cases, the matmul itself is not being tested so it's OK to run it with
  higher precision.

  If a matmul itself is being tested, or some other op which uses matmul, use
  `run_without_tensor_float_32` instead.

  This also casts complex64 inputs to complex128, since TensorFloat-32 can also
  be used with complex64

  Args:
    a: First input to tf.linalg.matmul
    b: Second input to tf.linalg.matmul
    args: Other positional arguments to tf.linalg.matmul
    **kwargs: Other keyword arguments to tf.linalg.matmul

  Returns:
    A tensor with the same type as `a`.
  """
  if config.tensor_float_32_execution_enabled() and a.dtype == "float32":
    a = math_ops.cast(a, "float64")
    b = math_ops.cast(b, "float64")
    ret = math_ops.matmul(a, b, *args, **kwargs)
    return math_ops.cast(ret, a.dtype)
  elif config.tensor_float_32_execution_enabled() and a.dtype == "complex64":
    a = math_ops.cast(a, "complex128")
    b = math_ops.cast(b, "complex128")
    ret = math_ops.matmul(a, b, *args, **kwargs)
    return math_ops.cast(ret, a.dtype)
  else:
    return math_ops.matmul(a, b, *args, **kwargs)


class EagerSessionWarner:

  def __getattr__(self, attr):
    raise AttributeError(
        "Trying to access properties or call methods on the result of "
        "self.session(), self.cached_session(), etc while eager execution "
        "is enabled. If you're porting this test case to TF 2.0, either "
        "adapt the test to work with eager execution or insert a call to "
        "tf.disable_eager_execution() in the main() function of this test "
        "file.")

# TODO(b/286583977): Set it to True and remove.
_ENABLE_AUTO_BOTH_MODES = False


@tf_export("test.TestCase")
class TensorFlowTestCase(googletest.TestCase):
  """Base class for tests that need to test TensorFlow."""

  def __init__(self, methodName="runTest"):  # pylint: disable=invalid-name
    super().__init__(methodName)
    # Make sure we get unfiltered stack traces during the test
    traceback_utils.disable_traceback_filtering()
    if is_xla_enabled():
      pywrap_tf_session.TF_SetXlaAutoJitMode("2")
      pywrap_tf_session.TF_SetXlaMinClusterSize(1)
      pywrap_tf_session.TF_SetXlaEnableLazyCompilation(False)
      pywrap_tf_session.TF_SetTfXlaCpuGlobalJit(True)
      # Constant folding secretly runs code on TF:Classic CPU, so we also
      # disable it here.
      pywrap_tf_session.TF_SetXlaConstantFoldingDisabled(True)

    # Check if the mlir bridge has been explicitly enabled or disabled. If
    # is_mlir_bridge_enabled() returns None, the user did not explictly enable
    # or disable the bridge so do not update enable_mlir_bridge.
    if is_mlir_bridge_enabled():
      context.context().enable_mlir_bridge = True
    elif is_mlir_bridge_enabled() is not None:
      context.context().enable_mlir_bridge = False

    self._threads = []
    self._tempdir = None
    self._cached_session = None
    self._test_start_time = None
    # This flag provides the ability to control whether the graph mode gets
    # initialized for TF1 or not. Initializing for TF1, which is what was
    # happening earlier, was preventing enablement of 'eager mode' in the test.
    self._set_default_seed = True

  def setUp(self):
    super().setUp()
    self._ClearCachedSession()
    random.seed(random_seed.DEFAULT_GRAPH_SEED)
    np.random.seed(random_seed.DEFAULT_GRAPH_SEED)
    # Note: The following line is necessary because some test methods may error
    # out from within nested graph contexts (e.g., via assertRaises and
    # assertRaisesRegexp), which may leave ops._default_graph_stack non-empty
    # under certain versions of Python. That would cause
    # ops.reset_default_graph() to throw an exception if the stack were not
    # cleared first.
    ops._default_graph_stack.reset()  # pylint: disable=protected-access
    ops.reset_default_graph()
    if self._set_default_seed:
      random_seed.set_random_seed(random_seed.DEFAULT_GRAPH_SEED)
    # Reset summary writer in case another test used set_as_default() with their
    # summary writer.
    summary_state = summary_ops_v2._summary_state  # pylint: disable=protected-access
    summary_state.writer = None

    # Avoiding calling setUp() for the poorly named test_session method.
    if self.id().endswith(".test_session"):
      self.skipTest("Not a test.")

    self._test_start_time = time.time()

  @classmethod
  def setUpClass(cls):
    super().setUpClass()
    if _ENABLE_AUTO_BOTH_MODES:
      run_all_in_graph_and_eager_modes(cls)

  def tearDown(self):
    # If a subclass overrides setUp and doesn't call the parent class's setUp,
    # then we may not have set the start time.
    if self._test_start_time is not None:
      logging.info("time(%s): %ss", self.id(),
                   round(time.time() - self._test_start_time, 2))

    for thread in self._threads:
      thread.check_termination()

    self._ClearCachedSession()
    super().tearDown()

  def _ClearCachedSession(self):
    if self._cached_session is not None:
      self._cached_session.close()
      self._cached_session = None

  def get_temp_dir(self) -> str:
    """Returns a unique temporary directory for the test to use.

    If you call this method multiple times during in a test, it will return the
    same folder. However, across different runs the directories will be
    different. This will ensure that across different runs tests will not be
    able to pollute each others environment.
    If you need multiple unique directories within a single test, you should
    use tempfile.mkdtemp as follows:
      tempfile.mkdtemp(dir=self.get_temp_dir()):

    Returns:
      string, the path to the unique temporary directory created for this test.
    """
    if not self._tempdir:
      self._tempdir = tempfile.mkdtemp(dir=googletest.GetTempDir())
    return self._tempdir

  @contextlib.contextmanager
  def captureWritesToStream(self, stream) -> Iterator[CapturedWrites]:
    """A context manager that captures the writes to a given stream.

    This context manager captures all writes to a given stream inside of a
    `CapturedWrites` object. When this context manager is created, it yields
    the `CapturedWrites` object. The captured contents can be accessed  by
    calling `.contents()` on the `CapturedWrites`.

    For this function to work, the stream must have a file descriptor that
    can be modified using `os.dup` and `os.dup2`, and the stream must support
    a `.flush()` method. The default python sys.stdout and sys.stderr are
    examples of this. Note that this does not work in Colab or Jupyter
    notebooks, because those use alternate stdout streams.

    Example:
    ```python
    class MyOperatorTest(test_util.TensorFlowTestCase):
      def testMyOperator(self):
        input = [1.0, 2.0, 3.0, 4.0, 5.0]
        with self.captureWritesToStream(sys.stdout) as captured:
          result = MyOperator(input).eval()
        self.assertStartsWith(captured.contents(), "This was printed.")
    ```

    Args:
      stream: The stream whose writes should be captured. This stream must have
        a file descriptor, support writing via using that file descriptor, and
        must have a `.flush()` method.

    Yields:
      A `CapturedWrites` object that contains all writes to the specified stream
      made during this context.
    """
    stream.flush()
    fd = stream.fileno()
    tmp_file, tmp_file_path = tempfile.mkstemp(dir=self.get_temp_dir())
    orig_fd = os.dup(fd)
    os.dup2(tmp_file, fd)
    try:
      yield CapturedWrites(tmp_file_path)
    finally:
      os.close(tmp_file)
      os.dup2(orig_fd, fd)

  def _AssertProtoEquals(self, a, b, msg=None, relative_tolerance=None):
    """Asserts that a and b are the same proto.

    Uses ProtoEq() first, as it returns correct results
    for floating point attributes, and then use assertProtoEqual()
    in case of failure as it provides good error messages.

    Args:
      a: a proto.
      b: another proto.
      msg: Optional message to report on failure.
      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.
    """
    if not compare.ProtoEq(a, b):
      compare.assertProtoEqual(
          self,
          a,
          b,
          normalize_numbers=True,
          msg=msg,
          relative_tolerance=relative_tolerance,
      )

  def assertProtoEquals(
      self,
      expected_message_maybe_ascii,
      message,
      msg=None,
      relative_tolerance=None,
  ):
    """Asserts that message is same as parsed expected_message_ascii.

    Creates another prototype of message, reads the ascii message into it and
    then compares them using self._AssertProtoEqual().

    Args:
      expected_message_maybe_ascii: proto message in original or ascii form.
      message: the message to validate.
      msg: Optional message to report on failure.
      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.
    """
    if isinstance(expected_message_maybe_ascii, type(message)):
      expected_message = expected_message_maybe_ascii
      self._AssertProtoEquals(
          expected_message,
          message,
          msg=msg,
          relative_tolerance=relative_tolerance,
      )
    elif isinstance(expected_message_maybe_ascii, (str, bytes)):
      expected_message = type(message)()
      text_format.Merge(
          expected_message_maybe_ascii,
          expected_message,
          descriptor_pool=descriptor_pool.Default())
      self._AssertProtoEquals(
          expected_message,
          message,
          msg=msg,
          relative_tolerance=relative_tolerance,
      )
    else:
      assert False, ("Can't compare protos of type %s and %s." %
                     (type(expected_message_maybe_ascii), type(message)))

  def assertProtoEqualsVersion(
      self,
      expected,
      actual,
      producer=versions.GRAPH_DEF_VERSION,
      min_consumer=versions.GRAPH_DEF_VERSION_MIN_CONSUMER,
      msg=None):
    expected = "versions { producer: %d min_consumer: %d };\n%s" % (
        producer, min_consumer, expected)
    self.assertProtoEquals(expected, actual, msg=msg)

  def assertStartsWith(self, actual, expected_start, msg=None):
    """Assert that actual.startswith(expected_start) is True.

    Args:
      actual: str
      expected_start: str
      msg: Optional message to report on failure.
    """
    if not actual.startswith(expected_start):
      fail_msg = "%r does not start with %r" % (actual, expected_start)
      fail_msg += " : %r" % (msg) if msg else ""
      self.fail(fail_msg)

  def _eval_tensor(self, tensor):
    if tensor is None:
      return None
    elif callable(tensor):
      return self._eval_helper(tensor())
    else:
      try:
        # for compatibility with TF1 test cases
        if sparse_tensor.is_sparse(tensor):
          return sparse_tensor.SparseTensorValue(tensor.indices.numpy(),
                                                 tensor.values.numpy(),
                                                 tensor.dense_shape.numpy())
        elif ragged_tensor.is_ragged(tensor):
          return ragged_tensor_value.RaggedTensorValue(
              self._eval_tensor(tensor.values),
              self._eval_tensor(tensor.row_splits))
        elif isinstance(tensor, indexed_slices.IndexedSlices):
          return indexed_slices.IndexedSlicesValue(
              values=tensor.values.numpy(),
              indices=tensor.indices.numpy(),
              dense_shape=None
              if tensor.dense_shape is None else tensor.dense_shape.numpy())
        else:
          if hasattr(tensor, "numpy") and callable(tensor.numpy):
            return tensor.numpy()
          else:
            # Try our best to convert CompositeTensor components to NumPy
            # arrays. Officially, we don't support NumPy arrays as
            # CompositeTensor components. So don't be surprised if this doesn't
            # work.
            return nest.map_structure(lambda t: t.numpy(), tensor,
                                      expand_composites=True)
      except AttributeError as e:
        raise ValueError(f"Unsupported type {type(tensor).__name__!r}.") from e

  def _eval_helper(self, tensors):
    if tensors is None:
      return None
    return nest.map_structure(self._eval_tensor, tensors)

  def evaluate(
      self, tensors
  ) -> Union[
      ragged_tensor_value.RaggedTensorValue,
      sparse_tensor.SparseTensorValue,
      None
  ]:
    """Evaluates tensors and returns numpy values.

    Args:
      tensors: A Tensor or a nested list/tuple of Tensors.

    Returns:
      tensors numpy values.
    """
    if context.executing_eagerly():
      return self._eval_helper(tensors)
    else:
      sess = ops.get_default_session()
      flattened_tensors = nest.flatten(tensors)
      if sess is None:
        with self.test_session() as sess:
          flattened_results = sess.run(flattened_tensors)
      else:
        flattened_results = sess.run(flattened_tensors)
      return nest.pack_sequence_as(tensors, flattened_results)

  # pylint: disable=g-doc-return-or-yield
  # pylint: disable=redefined-outer-name
  @contextlib.contextmanager
  def session(
      self,
      graph: Optional[ops.Graph] = None,
      config: Optional[config_pb2.ConfigProto] = None,
      use_gpu: bool = True,
      force_gpu: bool = False,
  ) -> Iterator[s.Session]:
    """A context manager for a TensorFlow Session for use in executing tests.

    Note that this will set this session and the graph as global defaults.

    Use the `use_gpu` and `force_gpu` options to control where ops are run. If
    `force_gpu` is True, all ops are pinned to `/device:GPU:0`. Otherwise, if
    `use_gpu` is True, TensorFlow tries to run as many ops on the GPU as
    possible. If both `force_gpu and `use_gpu` are False, all ops are pinned to
    the CPU.

    Example:

    ``` python
    class MyOperatorTest(test_util.TensorFlowTestCase):
      def testMyOperator(self):
        with self.session():
          valid_input = [1.0, 2.0, 3.0, 4.0, 5.0]
          result = MyOperator(valid_input).eval()
          self.assertEqual(result, [1.0, 2.0, 3.0, 5.0, 8.0]
          invalid_input = [-1.0, 2.0, 7.0]
          with self.assertRaisesOpError("negative input not supported"):
            MyOperator(invalid_input).eval()
    ```

    Args:
      graph: Optional graph to use during the returned session.
      config: An optional config_pb2.ConfigProto to use to configure the
        session.
      use_gpu: If True, attempt to run as many ops as possible on GPU.
      force_gpu: If True, pin all ops to `/device:GPU:0`.

    Yields:
      A Session object that should be used as a context manager to surround
      the graph building and execution code in a test case.
    """
    if context.executing_eagerly():
      yield EagerSessionWarner()
    else:
      with self._create_session(graph, config, force_gpu) as sess:
        with self._constrain_devices_and_set_default(sess, use_gpu, force_gpu):
          yield sess

  @contextlib.contextmanager
  def cached_session(
      self,
      graph: Optional[ops.Graph] = None,
      config: Optional[config_pb2.ConfigProto] = None,
      use_gpu: bool = True,
      force_gpu: bool = False,
  ) -> Iterator[s.Session]:
    """Returns a TensorFlow Session for use in executing tests.

    This method behaves differently than self.session(): for performance reasons
    `cached_session` will by default reuse the same session within the same
    test. The session returned by this function will only be closed at the end
    of the test (in the TearDown function).

    Use the `use_gpu` and `force_gpu` options to control where ops are run. If
    `force_gpu` is True, all ops are pinned to `/device:GPU:0`. Otherwise, if
    `use_gpu` is True, TensorFlow tries to run as many ops on the GPU as
    possible. If both `force_gpu and `use_gpu` are False, all ops are pinned to
    the CPU.

    Example:
    ```python
    class MyOperatorTest(test_util.TensorFlowTestCase):
      def testMyOperator(self):
        with self.cached_session() as sess:
          valid_input = [1.0, 2.0, 3.0, 4.0, 5.0]
          result = MyOperator(valid_input).eval()
          self.assertEqual(result, [1.0, 2.0, 3.0, 5.0, 8.0]
          invalid_input = [-1.0, 2.0, 7.0]
          with self.assertRaisesOpError("negative input not supported"):
            MyOperator(invalid_input).eval()
    ```

    Args:
      graph: Optional graph to use during the returned session.
      config: An optional config_pb2.ConfigProto to use to configure the
        session.
      use_gpu: If True, attempt to run as many ops as possible on GPU.
      force_gpu: If True, pin all ops to `/device:GPU:0`.

    Yields:
      A Session object that should be used as a context manager to surround
      the graph building and execution code in a test case.
    """
    if context.executing_eagerly():
      yield FakeEagerSession(self)
    else:
      sess = self._get_cached_session(
          graph, config, force_gpu, crash_if_inconsistent_args=True)
      with self._constrain_devices_and_set_default(sess, use_gpu,
                                                   force_gpu) as cached:
        yield cached

  @contextlib.contextmanager
  @deprecation.deprecated(None, "Use `self.session()` or "
                          "`self.cached_session()` instead.")
  def test_session(
      self,
      graph: Optional[ops.Graph] = None,
      config: Optional[config_pb2.ConfigProto] = None,
      use_gpu: bool = True,
      force_gpu: bool = False,
  ) -> Iterator[s.Session]:
    """Use cached_session instead."""
    if self.id().endswith(".test_session"):
      self.skipTest(
          "Tests that have the name \"test_session\" are automatically skipped "
          "by TensorFlow test fixture, as the name is reserved for creating "
          "sessions within tests. Please rename your test if you have a test "
          "with this name.")
    if context.executing_eagerly():
      yield None
    else:
      if graph is None:
        sess = self._get_cached_session(
            graph, config, force_gpu, crash_if_inconsistent_args=False)
        with self._constrain_devices_and_set_default(sess, use_gpu,
                                                     force_gpu) as cached:
          yield cached
      else:
        with self.session(graph, config, use_gpu, force_gpu) as sess:
          yield sess

  # pylint: enable=g-doc-return-or-yield

  class _CheckedThread(object):
    """A wrapper class for Thread that asserts successful completion.

    This class should be created using the TensorFlowTestCase.checkedThread()
    method.
    """

    def __init__(
        self,
        testcase: "TensorFlowTestCase",
        target: Callable[..., Any],
        args: Optional[tuple[Any, ...]] = None,
        kwargs: Optional[dict[str, Any]] = None,
    ):
      """Constructs a new instance of _CheckedThread.

      Args:
        testcase: The TensorFlowTestCase for which this thread is being created.
        target: A callable object representing the code to be executed in the
          thread.
        args: A tuple of positional arguments that will be passed to target.
        kwargs: A dictionary of keyword arguments that will be passed to target.
      """
      self._testcase = testcase
      self._target = target
      self._args: tuple[Any, ...] = () if args is None else args
      self._kwargs: dict[str, Any] = {} if kwargs is None else kwargs
      self._thread = threading.Thread(target=self._protected_run)
      self._exception = None

      self._is_thread_joined = False

    def _protected_run(self) -> None:
      """Target for the wrapper thread. Sets self._exception on failure."""
      try:
        self._target(*self._args, **self._kwargs)
      except Exception as e:  # pylint: disable=broad-except
        self._exception = e

    def start(self) -> None:
      """Starts the thread's activity.

      This must be called at most once per _CheckedThread object. It arranges
      for the object's target to be invoked in a separate thread of control.
      """
      self._thread.start()

    def join(self) -> None:
      """Blocks until the thread terminates.

      Raises:
        self._testcase.failureException: If the thread terminates with due to
          an exception.
      """
      self._is_thread_joined = True
      self._thread.join()
      if self._exception is not None:
        self._testcase.fail("Error in checkedThread: %s" % str(self._exception))

    def is_alive(self) -> bool:
      """Returns whether the thread is alive.

      This method returns True just before the run() method starts
      until just after the run() method terminates.

      Returns:
        True if the thread is alive, otherwise False.
      """
      return self._thread.is_alive()

    def check_termination(self) -> None:
      """Returns whether the checked thread was properly used and did terminate.

      Every checked thread should be "join"ed after starting, and before the
      test tears down. If it is not joined, it is possible the thread will hang
      and cause flaky failures in tests.

      Raises:
        self._testcase.failureException: If check_termination was called before
        thread was joined.

        RuntimeError: If the thread is not terminated. This means thread was not
        joined with the main thread.
      """
      if self._is_thread_joined:
        if self.is_alive():
          raise RuntimeError(
              "Thread was not joined with main thread, and is still running "
              "when the test finished.")
      else:
        self._testcase.fail("A checked thread was not joined.")

  def checkedThread(
      self,
      target: Callable[..., Any],
      args: Optional[tuple[Any, ...]] = None,
      kwargs: Optional[dict[str, Any]] = None,
  ) -> _CheckedThread:
    """Returns a Thread wrapper that asserts 'target' completes successfully.

    This method should be used to create all threads in test cases, as
    otherwise there is a risk that a thread will silently fail, and/or
    assertions made in the thread will not be respected.

    Args:
      target: A callable object to be executed in the thread.
      args: The argument tuple for the target invocation. Defaults to ().
      kwargs: A dictionary of keyword arguments for the target invocation.
        Defaults to {}.

    Returns:
      A wrapper for threading.Thread that supports start() and join() methods.
    """
    ret = TensorFlowTestCase._CheckedThread(self, target, args, kwargs)
    self._threads.append(ret)
    return ret

  # pylint: enable=invalid-name
  @py_func_if_in_function
  def assertNear(self, f1, f2, err, msg=None):
    """Asserts that two floats are near each other.

    Checks that |f1 - f2| < err and asserts a test failure
    if not.

    Args:
      f1: A float value.
      f2: A float value.
      err: A float value.
      msg: An optional string message to append to the failure message.
    """
    # f1 == f2 is needed here as we might have: f1, f2 = inf, inf
    self.assertTrue(
        f1 == f2 or math.fabs(f1 - f2) <= err, "%f != %f +/- %f%s" %
        (f1, f2, err, " (%s)" % msg if msg is not None else ""))

  @py_func_if_in_function
  def assertArrayNear(self, farray1, farray2, err, msg=None):
    """Asserts that two float arrays are near each other.

    Checks that for all elements of farray1 and farray2
    |f1 - f2| < err.  Asserts a test failure if not.

    Args:
      farray1: a list of float values.
      farray2: a list of float values.
      err: a float value.
      msg: Optional message to report on failure.
    """
    self.assertEqual(len(farray1), len(farray2), msg=msg)
    for f1, f2 in zip(farray1, farray2):
      self.assertNear(float(f1), float(f2), err, msg=msg)

  def _NDArrayNear(self, ndarray1, ndarray2, err):
    return np.linalg.norm(ndarray1 - ndarray2) < err

  @py_func_if_in_function
  def assertNDArrayNear(self, ndarray1, ndarray2, err, msg=None):
    """Asserts that two numpy arrays have near values.

    Args:
      ndarray1: a numpy ndarray.
      ndarray2: a numpy ndarray.
      err: a float. The maximum absolute difference allowed.
      msg: Optional message to report on failure.
    """
    self.assertTrue(self._NDArrayNear(ndarray1, ndarray2, err), msg=msg)

  def _GetNdArray(self, a):
    # If a is tensor-like then convert it to ndarray
    if tensor_util.is_tf_type(a):
      if isinstance(a, ops._EagerTensorBase):
        a = a.numpy()
      else:
        a = self.evaluate(a)
    if not isinstance(a, np.ndarray):
      try:
        return np.array(a)
      except ValueError as e:
        # TODO(b/264461299): NumPy 1.24 no longer infers dtype=object from
        # ragged sequences.
        # See:
        # https://numpy.org/neps/nep-0034-infer-dtype-is-object.html
        # Fixing this correctly requires clarifying the API contract of this
        # function with respect to ragged sequences and possibly updating all
        # users. As a backwards compatibility measure, if array
        # creation fails with an "inhomogeneous shape" error, try again with
        # an explicit dtype=object, which should restore the previous behavior.
        if "inhomogeneous shape" in str(e):
          return np.array(a, dtype=object)
        else:
          raise
    return a

  def evaluate_if_both_tensors(self, a, b):
    if (tensor_util.is_tf_type(a) and tensor_util.is_tf_type(b) and
        not isinstance(a, ops._EagerTensorBase) and
        not isinstance(b, ops._EagerTensorBase)):
      return self.evaluate((a, b))
    else:
      return (a, b)

  def _assertArrayLikeAllClose(self, a, b, rtol=1e-6, atol=1e-6, msg=None):
    (a, b) = self.evaluate_if_both_tensors(a, b)
    a = self._GetNdArray(a)
    b = self._GetNdArray(b)
    # When the array rank is small, print its contents. Numpy array printing is
    # implemented using inefficient recursion so prints can cause tests to
    # time out.
    if a.shape != b.shape and (b.ndim <= 3 or b.size < 500):
      shape_mismatch_msg = ("Shape mismatch: expected %s, got %s with contents "
                            "%s.") % (a.shape, b.shape, b)
    else:
      shape_mismatch_msg = "Shape mismatch: expected %s, got %s." % (a.shape,
                                                                     b.shape)
    self.assertEqual(a.shape, b.shape, shape_mismatch_msg)

    msgs = [msg]
    # np.allclose does not always work for our custom bfloat16 and float8
    # extension types when type promotions are involved, so we first cast any
    # arrays of such types to float32.
    a_dtype = a.dtype
    custom_dtypes = (dtypes.bfloat16.as_numpy_dtype,
                     dtypes.float8_e5m2.as_numpy_dtype,
                     dtypes.float8_e4m3fn.as_numpy_dtype)
    a = a.astype(np.float32) if a.dtype in custom_dtypes else a
    b = b.astype(np.float32) if b.dtype in custom_dtypes else b
    if not np.allclose(a, b, rtol=rtol, atol=atol):
      # Adds more details to np.testing.assert_allclose.
      #
      # NOTE: numpy.allclose (and numpy.testing.assert_allclose)
      # checks whether two arrays are element-wise equal within a
      # tolerance. The relative difference (rtol * abs(b)) and the
      # absolute difference atol are added together to compare against
      # the absolute difference between a and b.  Here, we want to
      # tell user which elements violate such conditions.
      cond = np.logical_or(
          np.abs(a - b) > atol + rtol * np.abs(b),
          np.isnan(a) != np.isnan(b))
      if a.ndim:
        x = a[np.where(cond)]
        y = b[np.where(cond)]
        msgs.append("not close where = {}".format(np.where(cond)))
      else:
        # np.where is broken for scalars
        x, y = a, b
      msgs.append("not close lhs = {}".format(x))
      msgs.append("not close rhs = {}".format(y))
      msgs.append("not close dif = {}".format(np.abs(x - y)))
      msgs.append("not close tol = {}".format(atol + rtol * np.abs(y)))
      msgs.append("dtype = {}, shape = {}".format(a_dtype, a.shape))
      # TODO(xpan): There seems to be a bug:
      # tensorflow/compiler/tests:binary_ops_test pass with float32
      # nan even though the equal_nan is False by default internally.
      np.testing.assert_allclose(
          a, b, rtol=rtol, atol=atol, err_msg="\n".join(msgs), equal_nan=True)

  def _assertAllCloseRecursive(self,
                               a,
                               b,
                               rtol=1e-6,
                               atol=1e-6,
                               path=None,
                               msg=None):
    if ragged_tensor.is_ragged(a) or ragged_tensor.is_ragged(b):
      return self._assertRaggedClose(a, b, rtol, atol, msg)
    path = path or []
    path_str = (("[" + "][".join(str(p) for p in path) + "]") if path else "")
    msg = msg if msg else ""

    # Check if a and/or b are namedtuples.
    if hasattr(a, "_asdict"):
      a = a._asdict()
    if hasattr(b, "_asdict"):
      b = b._asdict()
    a_is_dict = isinstance(a, collections_abc.Mapping)
    if a_is_dict != isinstance(b, collections_abc.Mapping):
      raise ValueError("Can't compare dict to non-dict, a%s vs b%s. %s" %
                       (path_str, path_str, msg))
    if a_is_dict:
      self.assertItemsEqual(
          a.keys(),
          b.keys(),
          msg="mismatched keys: a%s has keys %s, but b%s has keys %s. %s" %
          (path_str, a.keys(), path_str, b.keys(), msg))
      for k in a:
        path.append(k)
        self._assertAllCloseRecursive(
            a[k], b[k], rtol=rtol, atol=atol, path=path, msg=msg)
        del path[-1]
    elif isinstance(a, (list, tuple)):
      # Try to directly compare a, b as ndarrays; if not work, then traverse
      # through the sequence, which is more expensive.
      try:
        (a, b) = self.evaluate_if_both_tensors(a, b)
        a_as_ndarray = self._GetNdArray(a)
        b_as_ndarray = self._GetNdArray(b)
        self._assertArrayLikeAllClose(
            a_as_ndarray,
            b_as_ndarray,
            rtol=rtol,
            atol=atol,
            msg="Mismatched value: a%s is different from b%s. %s" %
            (path_str, path_str, msg))
      except (ValueError, TypeError, NotImplementedError) as e:
        if len(a) != len(b):
          raise ValueError(
              "Mismatched length: a%s has %d items, but b%s has %d items. %s" %
              (path_str, len(a), path_str, len(b), msg))
        for idx, (a_ele, b_ele) in enumerate(zip(a, b)):
          path.append(str(idx))
          self._assertAllCloseRecursive(
              a_ele, b_ele, rtol=rtol, atol=atol, path=path, msg=msg)
          del path[-1]
    # a and b are ndarray like objects
    else:
      try:
        self._assertArrayLikeAllClose(
            a,
            b,
            rtol=rtol,
            atol=atol,
            msg=("Mismatched value: a%s is different from b%s. %s" %
                 (path_str, path_str, msg)))
      except TypeError as e:
        msg = ("Error: a%s has %s, but b%s has %s. %s" %
               (path_str, type(a), path_str, type(b), msg))
        e.args = ((e.args[0] + " : " + msg,) + e.args[1:])
        raise

  @py_func_if_in_function
  def assertAllClose(self, a, b, rtol=1e-6, atol=1e-6, msg=None):
    """Asserts that two structures of numpy arrays or Tensors, have near values.

    `a` and `b` can be arbitrarily nested structures. A layer of a nested
    structure can be a `dict`, `namedtuple`, `tuple` or `list`.

    Note: the implementation follows
    [`numpy.allclose`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.allclose.html)
    (and numpy.testing.assert_allclose). It checks whether two arrays are
    element-wise equal within a tolerance. The relative difference
    (`rtol * abs(b)`) and the absolute difference `atol` are added together
    to compare against the absolute difference between `a` and `b`.

    Args:
      a: The expected numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor), or any arbitrarily nested of
        structure of these.
      b: The actual numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor), or any arbitrarily nested of
        structure of these.
      rtol: relative tolerance.
      atol: absolute tolerance.
      msg: Optional message to report on failure.

    Raises:
      ValueError: if only one of `a[p]` and `b[p]` is a dict or
          `a[p]` and `b[p]` have different length, where `[p]` denotes a path
          to the nested structure, e.g. given `a = [(1, 1), {'d': (6, 7)}]` and
          `[p] = [1]['d']`, then `a[p] = (6, 7)`.
    """
    self._assertAllCloseRecursive(a, b, rtol=rtol, atol=atol, msg=msg)

  @py_func_if_in_function
  def assertAllCloseAccordingToType(self,
                                    a,
                                    b,
                                    rtol=1e-6,
                                    atol=1e-6,
                                    float_rtol=1e-6,
                                    float_atol=1e-6,
                                    half_rtol=1e-3,
                                    half_atol=1e-3,
                                    bfloat16_rtol=1e-2,
                                    bfloat16_atol=1e-2,
                                    msg=None):
    """Like assertAllClose, but also suitable for comparing fp16 arrays.

    In particular, the tolerance is reduced to 1e-3 if at least
    one of the arguments is of type float16.

    Args:
      a: the expected numpy ndarray or anything can be converted to one.
      b: the actual numpy ndarray or anything can be converted to one.
      rtol: relative tolerance.
      atol: absolute tolerance.
      float_rtol: relative tolerance for float32.
      float_atol: absolute tolerance for float32.
      half_rtol: relative tolerance for float16.
      half_atol: absolute tolerance for float16.
      bfloat16_rtol: relative tolerance for bfloat16.
      bfloat16_atol: absolute tolerance for bfloat16.
      msg: Optional message to report on failure.
    """
    (a, b) = self.evaluate_if_both_tensors(a, b)
    a = self._GetNdArray(a)
    b = self._GetNdArray(b)
    # types with lower tol are put later to overwrite previous ones.
    if (a.dtype == np.float32 or b.dtype == np.float32 or
        a.dtype == np.complex64 or b.dtype == np.complex64):
      rtol = max(rtol, float_rtol)
      atol = max(atol, float_atol)
    if a.dtype == np.float16 or b.dtype == np.float16:
      rtol = max(rtol, half_rtol)
      atol = max(atol, half_atol)
    if (a.dtype == dtypes.bfloat16.as_numpy_dtype or
        b.dtype == dtypes.bfloat16.as_numpy_dtype):
      rtol = max(rtol, bfloat16_rtol)
      atol = max(atol, bfloat16_atol)

    self.assertAllClose(a, b, rtol=rtol, atol=atol, msg=msg)

  @py_func_if_in_function
  def assertNotAllClose(self, a, b, rtol=1e-6, atol=1e-6, msg=None):
    """Assert that two numpy arrays, or Tensors, do not have near values.

    Args:
      a: The expected numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor), or any arbitrarily nested of
        structure of these.
      b: The actual numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor), or any arbitrarily nested of
        structure of these.
      rtol: relative tolerance.
      atol: absolute tolerance.
      msg: Optional message to report on failure.

    Raises:
      AssertionError: If `a` and `b` are unexpectedly close at all elements.
    """
    try:
      self.assertAllClose(a, b, rtol=rtol, atol=atol, msg=msg)
    except AssertionError:
      return
    msg = msg or ""
    raise AssertionError("The two values are close at all elements. %s" % msg)

  @py_func_if_in_function
  def assertAllEqual(self, a, b, msg=None):
    """Asserts that two numpy arrays or Tensors have the same values.

    Args:
      a: the expected numpy ndarray or anything can be converted to one.
      b: the actual numpy ndarray or anything can be converted to one.
      msg: Optional message to report on failure.
    """
    if (ragged_tensor.is_ragged(a) or ragged_tensor.is_ragged(b)):
      return self._assertRaggedEqual(a, b, msg)
    msg = msg if msg else ""
    (a, b) = self.evaluate_if_both_tensors(a, b)
    a = self._GetNdArray(a)
    b = self._GetNdArray(b)
    # Arbitrary bounds so that we don't print giant tensors.
    if (b.ndim <= 3 or b.size < 500):
      self.assertEqual(
          a.shape, b.shape, "Shape mismatch: expected %s, got %s."
          " Contents: %r. \n%s." % (a.shape, b.shape, b, msg))
    else:
      self.assertEqual(
          a.shape, b.shape, "Shape mismatch: expected %s, got %s."
          " %s" % (a.shape, b.shape, msg))

    same = (a == b)

    if dtypes.as_dtype(a.dtype).is_floating:
      same = np.logical_or(same, np.logical_and(np.isnan(a), np.isnan(b)))
    msgs = [msg]
    if not np.all(same):
      # Adds more details to np.testing.assert_array_equal.
      diff = np.logical_not(same)
      if a.ndim:
        x = a[np.where(diff)]
        y = b[np.where(diff)]
        msgs.append("not equal where = {}".format(np.where(diff)))
      else:
        # np.where is broken for scalars
        x, y = a, b
      msgs.append("not equal lhs = %r" % x)
      msgs.append("not equal rhs = %r" % y)

      if (a.dtype.kind != b.dtype.kind and
          {a.dtype.kind, b.dtype.kind}.issubset({"U", "S", "O"})):
        a_list = []
        b_list = []
        # OK to flatten `a` and `b` because they are guaranteed to have the
        # same shape.
        for out_list, flat_arr in [(a_list, a.flat), (b_list, b.flat)]:
          for item in flat_arr:
            if isinstance(item, str):
              out_list.append(item.encode("utf-8"))
            else:
              out_list.append(item)
        a = np.array(a_list)
        b = np.array(b_list)

      np.testing.assert_array_equal(a, b, err_msg="\n".join(msgs))

  @py_func_if_in_function
  def assertNotAllEqual(self, a, b, msg=None):
    """Asserts that two numpy arrays or Tensors do not have the same values.

    Args:
      a: the expected numpy ndarray or anything can be converted to one.
      b: the actual numpy ndarray or anything can be converted to one.
      msg: Optional message to report on failure.
    """
    try:
      self.assertAllEqual(a, b)
    except AssertionError:
      return
    msg = msg or ""
    raise AssertionError("The two values are equal at all elements. %s" % msg)

  @py_func_if_in_function
  def assertAllGreater(self, a, comparison_target):
    """Assert element values are all greater than a target value.

    Args:
      a: The numpy `ndarray`, or anything that can be converted into a numpy
        `ndarray` (including Tensor).
      comparison_target: The target value of comparison.
    """
    (a, comparison_target) = self.evaluate_if_both_tensors(a, comparison_target)
    a = self._GetNdArray(a)
    self.assertGreater(np.min(a), comparison_target)

  @py_func_if_in_function
  def assertAllLess(self, a, comparison_target):
    """Assert element values are all less than a target value.

    Args:
      a: The numpy `ndarray`, or anything that can be converted into a numpy
        `ndarray` (including Tensor).
      comparison_target: The target value of comparison.
    """
    (a, comparison_target) = self.evaluate_if_both_tensors(a, comparison_target)
    a = self._GetNdArray(a)
    self.assertLess(np.max(a), comparison_target)

  @py_func_if_in_function
  def assertAllGreaterEqual(self, a, comparison_target):
    """Assert element values are all greater than or equal to a target value.

    Args:
      a: The numpy `ndarray`, or anything that can be converted into a numpy
        `ndarray` (including Tensor).
      comparison_target: The target value of comparison.
    """
    (a, comparison_target) = self.evaluate_if_both_tensors(a, comparison_target)
    a = self._GetNdArray(a)
    self.assertGreaterEqual(np.min(a), comparison_target)

  @py_func_if_in_function
  def assertAllLessEqual(self, a, comparison_target):
    """Assert element values are all less than or equal to a target value.

    Args:
      a: The numpy `ndarray`, or anything that can be converted into a numpy
        `ndarray` (including Tensor).
      comparison_target: The target value of comparison.
    """
    (a, comparison_target) = self.evaluate_if_both_tensors(a, comparison_target)
    a = self._GetNdArray(a)
    self.assertLessEqual(np.max(a), comparison_target)

  def _format_subscripts(self, subscripts, value, limit=10, indent=2):
    """Generate a summary of ndarray subscripts as a list of str.

    If limit == N, this method will print up to the first N subscripts on
    separate
    lines. A line of ellipses (...) will be appended at the end if the number of
    subscripts exceeds N.

    Args:
      subscripts: The tensor (np.ndarray) subscripts, of the same format as
        np.where()'s return value, i.e., a tuple of arrays with each array
        corresponding to a dimension. E.g., (array([1, 1]), array([0, 1])).
      value: (np.ndarray) value of the tensor.
      limit: (int) The maximum number of indices to print.
      indent: (int) Number of characters to indent at the beginning of each
        line.

    Returns:
      (list of str) the multi-line representation of the subscripts and values,
        potentially with omission at the end.
    """
    lines = []
    subscripts = np.transpose(subscripts)
    prefix = " " * indent
    if np.ndim(value) == 0:
      return [prefix + "[0] : " + str(value)]
    for subscript in itertools.islice(subscripts, limit):
      lines.append(prefix + str(subscript) + " : " +
                   str(value[tuple(subscript)]))
    if len(subscripts) > limit:
      lines.append(prefix + "...")
    return lines

  @py_func_if_in_function
  def assertAllInRange(self,
                       target,
                       lower_bound,
                       upper_bound,
                       open_lower_bound=False,
                       open_upper_bound=False):
    """Assert that elements in a Tensor are all in a given range.

    Args:
      target: The numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor).
      lower_bound: lower bound of the range
      upper_bound: upper bound of the range
      open_lower_bound: (`bool`) whether the lower bound is open (i.e., > rather
        than the default >=)
      open_upper_bound: (`bool`) whether the upper bound is open (i.e., < rather
        than the default <=)

    Raises:
      AssertionError:
        if the value tensor does not have an ordered numeric type (float* or
          int*), or
        if there are nan values, or
        if any of the elements do not fall in the specified range.
    """
    target = self._GetNdArray(target)
    if not (np.issubdtype(target.dtype, np.floating) or
            np.issubdtype(target.dtype, np.integer)):
      raise AssertionError(
          "The value of %s does not have an ordered numeric type, instead it "
          "has type: %s" % (target, target.dtype))

    nan_subscripts = np.where(np.isnan(target))
    if np.size(nan_subscripts):
      raise AssertionError(
          "%d of the %d element(s) are NaN. "
          "Subscripts(s) and value(s) of the NaN element(s):\n" %
          (len(nan_subscripts[0]), np.size(target)) +
          "\n".join(self._format_subscripts(nan_subscripts, target)))

    range_str = (("(" if open_lower_bound else "[") + str(lower_bound) + ", " +
                 str(upper_bound) + (")" if open_upper_bound else "]"))

    violations = (
        np.less_equal(target, lower_bound) if open_lower_bound else np.less(
            target, lower_bound))
    violations = np.logical_or(
        violations,
        np.greater_equal(target, upper_bound)
        if open_upper_bound else np.greater(target, upper_bound))
    violation_subscripts = np.where(violations)
    if np.size(violation_subscripts):
      raise AssertionError(
          "%d of the %d element(s) are outside the range %s. " %
          (len(violation_subscripts[0]), np.size(target), range_str) +
          "Subscript(s) and value(s) of the offending elements:\n" +
          "\n".join(self._format_subscripts(violation_subscripts, target)))

  @py_func_if_in_function
  def assertAllInSet(self, target, expected_set):
    """Assert that elements of a Tensor are all in a given closed set.

    Args:
      target: The numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor).
      expected_set: (`list`, `tuple` or `set`) The closed set that the elements
        of the value of `target` are expected to fall into.

    Raises:
      AssertionError:
        if any of the elements do not fall into `expected_set`.
    """
    target = self._GetNdArray(target)

    # Elements in target that are not in expected_set.
    diff = np.setdiff1d(target.flatten(), list(expected_set))
    if np.size(diff):
      raise AssertionError("%d unique element(s) are not in the set %s: %s" %
                           (np.size(diff), expected_set, diff))

  @py_func_if_in_function
  def assertDTypeEqual(self, target, expected_dtype):
    """Assert ndarray data type is equal to expected.

    Args:
      target: The numpy `ndarray`, or anything that can be converted into a
        numpy `ndarray` (including Tensor).
      expected_dtype: Expected data type.
    """
    target = self._GetNdArray(target)
    if not isinstance(target, list):
      arrays = [target]
    for arr in arrays:
      self.assertEqual(arr.dtype, expected_dtype)

  # pylint: disable=g-doc-return-or-yield
  @contextlib.contextmanager
  def assertRaisesWithPredicateMatch(self, exception_type,
                                     expected_err_re_or_predicate):
    """Returns a context manager to enclose code expected to raise an exception.

    If the exception is an OpError, the op stack is also included in the message
    predicate search.

    Args:
      exception_type: The expected type of exception that should be raised.
      expected_err_re_or_predicate: If this is callable, it should be a function
        of one argument that inspects the passed-in exception and returns True
        (success) or False (please fail the test). Otherwise, the error message
        is expected to match this regular expression partially.

    Returns:
      A context manager to surround code that is expected to raise an
      exception.
    """
    if callable(expected_err_re_or_predicate):
      predicate = expected_err_re_or_predicate
    else:

      def predicate(e):
        if isinstance(e, errors.OpError):
          e = cast(errors.OpError, e)
          err_str = cast(str, e.message)
          op = e.op
        else:
          err_str = str(e)
          op = None
        while op is not None:
          err_str += "\nCaused by: " + op.name
          op = op._original_op  # pylint: disable=protected-access
        logging.info("Searching within error strings: '%s' within '%s'",
                     expected_err_re_or_predicate, err_str)
        return re.search(expected_err_re_or_predicate, err_str)

    try:
      yield
      self.fail(exception_type.__name__ + " not raised")
    except Exception as e:  # pylint: disable=broad-except
      if not isinstance(e, exception_type) or not predicate(e):
        raise AssertionError("Exception of type %s: %s" %
                             (str(type(e)), str(e)))

  # pylint: enable=g-doc-return-or-yield

  def assertRaisesOpError(self, expected_err_re_or_predicate):
    return self.assertRaisesWithPredicateMatch(errors.OpError,
                                               expected_err_re_or_predicate)

  def assertRaisesIncompatibleShapesError(
      self, exception_type=errors.InvalidArgumentError):
    return self.assertRaisesWithPredicateMatch(
        exception_type, r"Incompatible shapes|Dimensions must be equal|"
        r"required broadcastable shapes")

  def assertShapeEqual(self, input_a, input_b, msg=None):
    """Asserts that two Numpy or TensorFlow objects have the same shape.

    For Tensors, this compares statically known shapes at compile time, not
    dynamic shapes at runtime.

    Args:
      input_a: A Numpy ndarray, Numpy scalar, or a Tensor.
      input_b: A Numpy ndarray, Numpy scalar, or a Tensor.
      msg: Optional message to report on failure.

    Raises:
      TypeError: If the arguments have the wrong type.
    """
    if not isinstance(input_a, (np.ndarray, np.generic, tensor_lib.Tensor)):
      raise TypeError(
          "input_a must be a Numpy ndarray, Numpy scalar, or a Tensor."
          f"Instead received {type(input_a)}")
    if not isinstance(input_b, (np.ndarray, np.generic, tensor_lib.Tensor)):
      raise TypeError(
          "input_b must be a Numpy ndarray, Numpy scalar, or a Tensor."
          f"Instead received {type(input_b)}")
    shape_a = input_a.get_shape().as_list() if isinstance(
        input_a, tensor_lib.Tensor) else input_a.shape
    shape_b = input_b.get_shape().as_list() if isinstance(
        input_b, tensor_lib.Tensor) else input_b.shape
    self.assertAllEqual(shape_a, shape_b, msg=msg)

  def assertDeviceEqual(self, device1, device2, msg=None):
    """Asserts that the two given devices are the same.

    Args:
      device1: A string device name or TensorFlow `DeviceSpec` object.
      device2: A string device name or TensorFlow `DeviceSpec` object.
      msg: Optional message to report on failure.
    """
    device1 = pydev.canonical_name(device1)
    device2 = pydev.canonical_name(device2)
    self.assertEqual(
        device1, device2,
        "Devices %s and %s are not equal. %s" % (device1, device2, msg))

  @py_func_if_in_function
  def assertDictEqual(self, a, b, msg=None):
    """Assert that two given dictionary of tensors are the same.

    Args:
      a: Expected dictionary with numpy ndarray or anything else that can be
        converted to one as values.
      b: Actual dictionary with numpy ndarray or anything else that can be
        converted to one as values.
      msg: Optional message to report on failure.
    """
    # To keep backwards compatibility, we first try the base class
    # assertDictEqual. If that fails we try the tensorflow one.
    try:
      super().assertDictEqual(a, b, msg)
    except Exception:  # pylint: disable=broad-except
      self.assertSameElements(a.keys(), b.keys())  # pylint: disable=g-assert-in-except
      for k, v in a.items():
        (a_k, b_k) = self.evaluate_if_both_tensors(v, b[k])
        a_k = self._GetNdArray(a_k)
        b_k = self._GetNdArray(b_k)
        if np.issubdtype(a_k.dtype, np.floating):
          self.assertAllClose(v, b[k], msg=k)
        else:
          self.assertAllEqual(v, b[k], msg=k)

  def _GetPyList(self, a):
    """Converts `a` to a nested python list."""
    if isinstance(a, ragged_tensor.RaggedTensor):
      a = cast(ragged_tensor_value.RaggedTensorValue, self.evaluate(a))
      return a.to_list()
    elif isinstance(a, tensor_lib.Tensor):
      a = self.evaluate(a)
      return a.tolist() if isinstance(a, np.ndarray) else a
    elif isinstance(a, np.ndarray):
      return a.tolist()
    elif isinstance(a, ragged_tensor_value.RaggedTensorValue):
      return a.to_list()
    else:
      return np.array(a, dtype=object).tolist()

  def _assertRaggedEqual(self, a, b, msg):
    """Asserts that two ragged tensors are equal."""
    a_list = self._GetPyList(a)
    b_list = self._GetPyList(b)
    self.assertEqual(a_list, b_list, msg)

    if not (isinstance(a, (list, tuple)) or isinstance(b, (list, tuple))):
      a_ragged_rank = a.ragged_rank if ragged_tensor.is_ragged(a) else 0
      b_ragged_rank = b.ragged_rank if ragged_tensor.is_ragged(b) else 0
      self.assertEqual(a_ragged_rank, b_ragged_rank, msg)

  def _assertRaggedClose(self, a, b, rtol, atol, msg=None):
    a_list = self._GetPyList(a)
    b_list = self._GetPyList(b)
    self._assertListCloseRecursive(a_list, b_list, rtol, atol, msg)

    if not (isinstance(a, (list, tuple)) or isinstance(b, (list, tuple))):
      a_ragged_rank = a.ragged_rank if ragged_tensor.is_ragged(a) else 0
      b_ragged_rank = b.ragged_rank if ragged_tensor.is_ragged(b) else 0
      self.assertEqual(a_ragged_rank, b_ragged_rank, msg)

  def _assertListCloseRecursive(self, a, b, rtol, atol, msg, path="value"):
    self.assertEqual(type(a), type(b))
    if isinstance(a, (list, tuple)):
      self.assertLen(a, len(b), "Length differs for %s" % path)
      for i in range(len(a)):
        self._assertListCloseRecursive(a[i], b[i], rtol, atol, msg,
                                       "%s[%s]" % (path, i))
    else:
      self._assertAllCloseRecursive(a, b, rtol, atol, path, msg)

  # Fix Python 3+ compatibility issues
  # pylint: disable=invalid-name

  # Silence a deprecation warning
  assertRaisesRegexp = googletest.TestCase.assertRaisesRegex

  # assertItemsEqual is assertCountEqual as of 3.2.
  assertItemsEqual = googletest.TestCase.assertCountEqual

  # pylint: enable=invalid-name

  @contextlib.contextmanager
  def _constrain_devices_and_set_default(
      self, sess: s.Session, use_gpu: bool, force_gpu: bool,
  ) -> Iterator[s.Session]:
    """Set the session and its graph to global default and constrain devices."""
    if context.executing_eagerly():
      yield None
    else:
      with sess.graph.as_default(), sess.as_default():
        if force_gpu:
          # Use the name of an actual device if one is detected, or
          # '/device:GPU:0' otherwise
          gpu_name = gpu_device_name()
          if not gpu_name:
            gpu_name = "/device:GPU:0"
          with sess.graph.device(gpu_name):
            yield sess
        elif use_gpu:
          yield sess
        else:
          with sess.graph.device("/device:CPU:0"):
            yield sess

  def _create_session(
      self,
      graph: Optional[ops.Graph],
      config: Optional[config_pb2.ConfigProto],
      force_gpu: bool,
  ) -> s.Session:
    """See session() for details."""

    def prepare_config(
        config: Optional[config_pb2.ConfigProto],
    ) -> config_pb2.ConfigProto:
      """Returns a config for sessions.

      Args:
        config: An optional config_pb2.ConfigProto to use to configure the
          session.

      Returns:
        A config_pb2.ConfigProto object.
      """
      # TODO(b/114333779): Enforce allow_soft_placement=False when
      # use_gpu=False. Currently many tests rely on the fact that any device
      # will be used even when a specific device is supposed to be used.
      allow_soft_placement = not force_gpu
      if config is None:
        config = context.context().config
        config.allow_soft_placement = allow_soft_placement
      elif not allow_soft_placement and config.allow_soft_placement:
        config_copy = context.context().config
        config = config_copy
        config.allow_soft_placement = False
      # Don't perform optimizations for tests so we don't inadvertently run
      # gpu ops on cpu
      config.graph_options.optimizer_options.opt_level = -1
      # Disable Grappler constant folding since some tests & benchmarks
      # use constant input and become meaningless after constant folding.
      # DO NOT DISABLE GRAPPLER OPTIMIZERS WITHOUT CONSULTING WITH THE
      # GRAPPLER TEAM.
      config.graph_options.rewrite_options.constant_folding = (
          rewriter_config_pb2.RewriterConfig.OFF)
      config.graph_options.rewrite_options.pin_to_host_optimization = (
          rewriter_config_pb2.RewriterConfig.OFF)
      return config

    return ErrorLoggingSession(graph=graph, config=prepare_config(config))

  def _get_cached_session(
      self,
      graph: Optional[ops.Graph] = None,
      config: Optional[config_pb2.ConfigProto] = None,
      force_gpu: bool = False,
      crash_if_inconsistent_args: bool = True,
  ) -> s.Session:
    """See cached_session() for documentation."""
    if self._cached_session is None:
      sess = self._create_session(
          graph=graph, config=config, force_gpu=force_gpu)
      self._cached_session = sess
      self._cached_graph = graph
      self._cached_config = config
      self._cached_force_gpu = force_gpu
      return sess
    else:
      if crash_if_inconsistent_args and self._cached_graph is not graph:
        raise ValueError("The graph used to get the cached session is "
                         "different than the one that was used to create the "
                         "session. Maybe create a new session with "
                         "self.session()")
      if crash_if_inconsistent_args and self._cached_config is not config:
        raise ValueError("The config used to get the cached session is "
                         "different than the one that was used to create the "
                         "session. Maybe create a new session with "
                         "self.session()")
      if crash_if_inconsistent_args and (self._cached_force_gpu is
                                         not force_gpu):
        raise ValueError(
            "The force_gpu value used to get the cached session is "
            "different than the one that was used to create the "
            "session. Maybe create a new session with "
            "self.session()")
      return self._cached_session


ASSIGNED_PORTS: set[int] = set()
lock = threading.Lock()


def pick_unused_port():
  """Returns an unused and unassigned local port."""
  import portpicker  # pylint: disable=g-import-not-at-top

  global ASSIGNED_PORTS
  with lock:
    while True:
      try:
        port = portpicker.pick_unused_port()
      except portpicker.NoFreePortFoundError as porterror:
        raise unittest.SkipTest("Flakes in portpicker library do not represent"
                                " TensorFlow errors.") from porterror
      if port > 10000 and port not in ASSIGNED_PORTS:
        ASSIGNED_PORTS.add(port)
        logging.info("Using local port %r", port)
        return port


@tf_export("test.create_local_cluster")
def create_local_cluster(
    num_workers: int,
    num_ps: int,
    protocol: str = "grpc",
    worker_config: Optional[config_pb2.ConfigProto] = None,
    ps_config: Optional[config_pb2.ConfigProto] = None,
) -> tuple[list[server_lib.Server], list[server_lib.Server]]:
  """Create and start local servers and return the associated `Server` objects.

  "PS" stands for "parameter server": a task responsible for storing and
  updating the model's parameters. Other tasks send updates to these parameters
  as they work on optimizing the parameters. This particular division of labor
  between tasks is not required, but is common for distributed training.

  Read more at https://www.tensorflow.org/guide/extend/architecture

  ![components](https://www.tensorflow.org/images/diag1.svg "components")


  Figure illustrates the interaction of these components.
  "/job:worker/task:0" and "/job:ps/task:0" are both tasks with worker services.


  Example:
  ```python
  workers, _ = tf.test.create_local_cluster(num_workers=2, num_ps=2)

  worker_sessions = [tf.compat.v1.Session(w.target) for w in workers]

  with tf.device("/job:ps/task:0"):
    ...
  with tf.device("/job:ps/task:1"):
    ...
  with tf.device("/job:worker/task:0"):
    ...
  with tf.device("/job:worker/task:1"):
    ...

  worker_sessions[0].run(...)
  ```

  Args:
    num_workers: Number of worker servers to start.
    num_ps: Number of PS servers to start.
    protocol: Communication protocol. Allowed values are documented in the
      documentation of `tf.distribute.Server`.
    worker_config: (optional) `tf.ConfigProto` to initialize workers. Can be
      used to instantiate multiple devices etc.
    ps_config: (optional) `tf.ConfigProto` to initialize PS servers.

  Returns:
    A tuple `(worker_servers, ps_servers)`.  `worker_servers` is a list
    of `num_workers` objects of type `tf.distribute.Server` (all running
    locally);
    and `ps_servers` is a list of `num_ps` objects of similar type.

  Raises:
    ImportError: if portpicker module was not found at load time
  """
  worker_ports = [pick_unused_port() for _ in range(num_workers)]
  ps_ports = [pick_unused_port() for _ in range(num_ps)]
  cluster_dict = {
      "worker": ["localhost:%s" % port for port in worker_ports],
      "ps": ["localhost:%s" % port for port in ps_ports]
  }
  cs = server_lib.ClusterSpec(cluster_dict)

  workers = [
      server_lib.Server(
          cs,
          job_name="worker",
          protocol=protocol,
          task_index=ix,
          config=worker_config,
          start=True) for ix in range(num_workers)
  ]
  ps_servers = [
      server_lib.Server(
          cs,
          job_name="ps",
          protocol=protocol,
          task_index=ix,
          config=ps_config,
          start=True) for ix in range(num_ps)
  ]

  return workers, ps_servers


def get_node_def_from_graph(
    node_name: str, graph_def: graph_pb2.GraphDef,
) -> Optional[node_def_pb2.NodeDef]:
  """Returns the `NodeDef` instance for given node name in the graph def.

  This method explores only the NodeDefs in `graph_def.node`.

  Args:
    node_name: Name of the NodeDef to search for.
    graph_def: An instance of `GraphDef` proto.

  Returns:
    the `NodeDef` instance whose name field matches the given node_name or None.
  """
  for node_def in graph_def.node:
    if node_def.name == node_name:
      return node_def
  return None


def set_producer_version(graph: ops.Graph, producer_version: int) -> None:
  """Sets graph.graph_def_versions.producer to `producer_version`."""
  # The C API doesn't expose altering GraphDefVersions. We can indirectly set
  # it via import_graph_def though.
  graph_def = graph_pb2.GraphDef()
  graph_def.versions.producer = producer_version
  with graph.as_default():
    importer.import_graph_def(graph_def)
  assert graph.graph_def_versions.producer, producer_version


@contextlib.contextmanager
def _fake_gradient_tape_context_manager():
  """tf.gradients(...) implemented as tf.GradientTape context manager interface.

  This is useful to test tf.gradients() in tests that uses tf.GradientTape().

  Yields:
    gradient tape instance that's implemented by tf.gradients() underneath.
  """
  try:
    class FakeGradientTape:

      def watch(self, x):
        pass

      def gradient(self, y, x, grad_ys=None):
        result = gradients_impl.gradients(y, x, grad_ys)

        # Unlike `tape.gradient()`, `tf.gradients()` returns a list for a single
        # element. So unpack if needed to match `tape.gradient()` behavior.
        if not isinstance(x, (list, tuple)):
          assert len(result) == 1
          return result[0]

        return result

    yield FakeGradientTape()
  finally:
    pass


class AbstractGradientTape:
  """Abstract GradientTape context manager that has multiple implementations.

  This is useful to test both tf.GradientTape() and tf.gradients() without
  duplicating tests.
  """

  def __init__(self, use_tape, persistent=False):
    self._use_tape = use_tape
    self._persistent = persistent

  def __enter__(self) -> backprop.GradientTape:
    if self._use_tape:
      self._tape_impl = backprop.GradientTape(persistent=self._persistent)
    else:
      self._tape_impl = _fake_gradient_tape_context_manager()
    return self._tape_impl.__enter__()

  def __exit__(self, exc_type, exc_val, exc_tb):
    self._tape_impl.__exit__(exc_type, exc_val, exc_tb)


@contextlib.contextmanager
def run_functions_eagerly(run_eagerly: bool) -> Iterator[None]:
  """Runs functions eagerly if `run_eagerly` is true.

  WARNING: Setting `run_eagerly` to True in tests running in V1 graph mode
  *WILL NOT* make the tf.function to run eagerly because eager is disabled by
  default in V1. Instead, tf.function will run as a traced graph function.

  Ensures that the state (for running functions eagerly) is back to the initial
  `def_function.RUN_FUNCTIONS_EAGERLY` state.

  Args:
    run_eagerly: Boolean determining whether to run the function eagerly or not.

  Raises:
    ValueError if `run_eagerly` is not a boolean.

  Yields:
    Nothing.
  """
  if not isinstance(run_eagerly, bool):
    raise ValueError(
        "Expected bool for `run_eagerly` but got {}".format(run_eagerly))

  is_eager = context.executing_eagerly()
  if not is_eager and run_eagerly:
    logging.warning(
        "Running tf.function eagerly in V1 graph mode is not supported. "
        "tf.function will be run as a traced graph function.")

  initial_state = def_function.functions_run_eagerly()
  def_function.run_functions_eagerly(run_eagerly)
  try:
    yield
  finally:
    def_function.run_functions_eagerly(initial_state)


class TestDelta:
  """A utility class to track increments to test counters."""

  def __init__(self, name, label):
    self.name = name
    self.label = label
    self.Reset()

  def Reset(self) -> None:
    self.last_value = _test_metrics_util.test_counter_value(
        self.name, self.label)

  def Get(self) -> int:
    value = _test_metrics_util.test_counter_value(self.name, self.label)
    return value - self.last_value


@tf_export("test.experimental.sync_devices")
def sync_devices() -> None:
  """Synchronizes all devices.

  By default, GPUs run asynchronously. This means that when you run an op on the
  GPU, like `tf.linalg.matmul`, the op may still be running on the GPU when the
  function returns. Non-GPU devices can also be made to run asynchronously by
  calling `tf.config.experimental.set_synchronous_execution(False)`. Calling
  `sync_devices()` blocks until pending ops have finished executing. This is
  primarily useful for measuring performance during a benchmark.

  For example, here is how you can measure how long `tf.linalg.matmul` runs:

  >>> import time
  >>> x = tf.random.normal((4096, 4096))
  >>> tf.linalg.matmul(x, x)  # Warmup.
  >>> tf.test.experimental.sync_devices()  # Block until warmup has completed.
  >>>
  >>> start = time.time()
  >>> y = tf.linalg.matmul(x, x)
  >>> tf.test.experimental.sync_devices()  # Block until matmul has completed.
  >>> end = time.time()
  >>> print(f'Time taken: {end - start}')

  If the call to `sync_devices()` was omitted, the time printed could be too
  small. This is because the op could still be running asynchronously when
  the line `end = time.time()` is executed.

  Raises:
    RuntimeError: If run outside Eager mode. This must be called in Eager mode,
      outside any `tf.function`s.
  """
  if not context.executing_eagerly():
    raise RuntimeError(
        "sync_devices() must only be called in Eager mode, outside tf.functions"
    )

  # There are two sources of asynchrony in TensorFlow:
  #
  # 1. On GPUs, kernels are run on a CUDA stream, which is inherently
  #    asynchronous.
  # 2. Calling `tf.config.experimental.set_synchronous_execution(False)` makes
  #    all ops asynchronous, in which case TensorFlow maintains internal queues
  #    of pending ops.
  #
  # Calling SyncDevice addresses source (1). Calling async_await addresses
  # source (2). It is important that SyncDevice() is called before async_wait(),
  # otherwise the SyncDevice op itself may still be pending on an internal
  # TensorFlow queue when the sync_devices() Python function returns.
  devices = config.list_logical_devices()
  for dev in devices:
    with ops.device(dev.name):
      gen_sync_ops.SyncDevice()
  context.async_wait()