chore: import upstream snapshot with attribution

tensorflow/python/compiler/BUILD

@@ -0,0 +1,16 @@
+# Description:
+# Python APIs for various Tensorflow backends.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "compiler",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+)

tensorflow/python/compiler/mlir/BUILD

@@ -0,0 +1,45 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "mlir",
+    srcs = ["mlir.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python:pywrap_mlir",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "mlir_test",
+    srcs = ["mlir_test.py"],
+    data = [
+        "multi_add.tflite",
+    ],
+    strict_deps = True,
+    tags = [
+        "no_mac",  # TODO(b/291995580): Reenable.
+        "no_pip",
+    ],
+    deps = [
+        ":mlir",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python:pywrap_mlir",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:logging_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+    ],
+)

tensorflow/python/compiler/mlir/mlir.py

@@ -0,0 +1,206 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""mlir is an experimental library that provides support APIs for MLIR."""
+
+from tensorflow.python import pywrap_mlir
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('mlir.experimental.convert_graph_def')
+def convert_graph_def(
+    graph_def, pass_pipeline='tf-standard-pipeline', show_debug_info=False
+):
+  """Import a GraphDef and convert it to a textual MLIR module.
+
+  This API is only intended for inspecting the internals of TensorFlow and the
+  string returned is at the moment intended for debugging purposes.
+
+  Args:
+    graph_def: An object of type graph_pb2.GraphDef or a textual proto
+      representation of a valid GraphDef.
+    pass_pipeline: A textual description of an MLIR Pass Pipeline to run on the
+      module, see MLIR documentation for the [textual pass pipeline
+      syntax](https://mlir.llvm.org/docs/PassManagement/#textual-pass-pipeline-specification).
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the graphdef.
+
+  Raises:
+    InvalidArgumentError: if graph_def is invalid or cannot be converted to
+      MLIR.
+  """
+  return pywrap_mlir.import_graphdef(graph_def, pass_pipeline, show_debug_info)
+
+
+@tf_export('mlir.experimental.convert_function')
+def convert_function(
+    concrete_function,
+    pass_pipeline='tf-standard-pipeline',
+    show_debug_info=False,
+):
+  """Import a ConcreteFunction and convert it to a textual MLIR module.
+
+  This API is only intended for inspecting the internals of TensorFlow and the
+  string returned is at the moment intended for debugging purposes.
+
+  A [tf.function](https://www.tensorflow.org/api_docs/python/tf/function) can be
+  imported and converted from TensorFlow to TensorFlow MLIR with this API by
+  extracting its ConcreteFunction (eagerly-executing wrapper around a
+  [tf.Graph](https://www.tensorflow.org/api_docs/python/tf/Graph)).
+
+  For example:
+  >>> @tf.function
+  ... def add(a, b):
+  ...   return a + b
+
+  >>> concrete_function = add.get_concrete_function(
+  ...     tf.TensorSpec(None, tf.dtypes.float32),
+  ...     tf.TensorSpec(None, tf.dtypes.float32))
+  >>> tf.mlir.experimental.convert_function(concrete_function)
+  '...module attributes {...} {...}...'
+
+  Args:
+    concrete_function: An object of type ConcreteFunction.
+    pass_pipeline: A textual description of an MLIR Pass Pipeline to run on the
+      module, see MLIR documentation for the [textual pass pipeline
+      syntax](https://mlir.llvm.org/docs/PassManagement/#textual-pass-pipeline-specification).
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    ConcreteFunction.
+
+  Raises:
+    InvalidArgumentError: if concrete_function is invalid or cannot be converted
+      to MLIR.
+  """
+  return pywrap_mlir.import_function(
+      concrete_function, pass_pipeline, show_debug_info
+  )
+
+
+@tf_export('mlir.experimental.convert_saved_model')
+def convert_saved_model(
+    saved_model_path, exported_names, show_debug_info=False
+):
+  """Converts a SavedModel to MLIR module.
+
+  Args:
+    saved_model_path: Path to SavedModel.
+    exported_names: Names to export.
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    SavedModel.
+  """
+  return pywrap_mlir.experimental_convert_saved_model_to_mlir(
+      saved_model_path, exported_names, show_debug_info
+  )
+
+
+@tf_export('mlir.experimental.convert_saved_model_v1')
+def convert_saved_model_v1(
+    saved_model_path,
+    exported_names,
+    tags,
+    lift_variables,
+    include_variables_in_initializers,
+    upgrade_legacy=True,
+    show_debug_info=False,
+):
+  """Converts a v1 SavedModel to MLIR module.
+
+  Args:
+    saved_model_path: Path to SavedModel.
+    exported_names: Names to export.
+    tags: MetaGraphDef to be loaded is identified by the supplied tags.
+    lift_variables: Whether to promote tf.VarHandleOp to resource arguments.
+    include_variables_in_initializers: Keeps the variables in initializers
+      before lifting variables.
+    upgrade_legacy: Functionalize the input graph before importing.
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    SavedModule.
+  """
+  return pywrap_mlir.experimental_convert_saved_model_v1_to_mlir(
+      saved_model_path,
+      exported_names,
+      tags,
+      lift_variables,
+      include_variables_in_initializers,
+      upgrade_legacy,
+      show_debug_info,
+  )
+
+
+@tf_export('mlir.experimental.run_pass_pipeline')
+def run_pass_pipeline(mlir_txt, pass_pipeline, show_debug_info=False):
+  """Runs a pipeline over input module.
+
+  Args:
+    mlir_txt: Textual representation of the MLIR module.
+    pass_pipeline: Pass pipeline to run on module.
+    show_debug_info: Whether to include locations in the emitted textual form.
+
+  Returns:
+    A textual representation of the MLIR module corresponding to the
+    transformed module.
+  """
+  return pywrap_mlir.experimental_run_pass_pipeline(
+      mlir_txt, pass_pipeline, show_debug_info
+  )
+
+
+@tf_export('mlir.experimental.write_bytecode')
+def experimental_write_bytecode(filename, mlir_txt):
+  """Writes an MLIR module out as bytecode.
+
+  Args:
+    filename: The filename to write to.
+    mlir_txt: The MLIR module in textual format.
+  """
+  pywrap_mlir.experimental_write_bytecode(filename, mlir_txt)
+
+
+@tf_export('mlir.experimental.tflite_to_tosa_bytecode')
+def tflite_to_tosa_bytecode(
+    flatbuffer,
+    bytecode,
+    use_external_constant=False,
+    ordered_input_arrays=None,
+    ordered_output_arrays=None,
+):
+  """Converts TFLite flatbuffer to TOSA dialect in MLIR bytecode.
+
+  Args:
+    flatbuffer: Path to flatbuffer.
+    bytecode: Path to output bytecode.
+    use_external_constant: Whether to create `tfl.external_const` instead of
+      `tfl.const`.
+    ordered_input_arrays:
+    ordered_output_arrays: If ordered_output_arrays is not empty, then the
+      function will only return nodes in ordered_output_arrays in the same order
+  """
+  pywrap_mlir.experimental_tflite_to_tosa_bytecode(
+      flatbuffer,
+      bytecode,
+      use_external_constant,
+      ordered_input_arrays,
+      ordered_output_arrays,
+  )

tensorflow/python/compiler/mlir/mlir_test.py

@@ -0,0 +1,162 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Tests for python.compiler.mlir."""
+
+from tensorflow.python.compiler.mlir import mlir
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.pywrap_mlir import import_graphdef
+
+
+class MLIRGraphDefImportTest(test.TestCase):
+
+  def testImport(self):
+    """Tests the basic flow of `tf.mlir.experimental.convert_graph_def`."""
+    mlir_module = mlir.convert_graph_def('')
+    # An empty graph should contain at least an empty main function.
+    self.assertIn('func @main', mlir_module)
+
+  def testInvalidPbtxt(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not parse input proto'):
+      mlir.convert_graph_def('some invalid proto')
+
+  def testGraphDefToTf(self):
+    """Tests the basic flow of `tf.mlir.experimental.convert_graph_def`
+
+        with tf-standard-pipeline converting all the way to the TF dialect.
+    """
+
+    tensor_shape = (10, 10)
+
+    @def_function.function(
+        input_signature=(
+            tensor_spec.TensorSpec(shape=tensor_shape, dtype=dtypes.float32),
+            tensor_spec.TensorSpec(shape=tensor_shape, dtype=dtypes.float32),
+        ))
+    def add_func(lhs, rhs):
+      return math_ops.add(lhs, rhs)
+
+    tf_graph_def = add_func.get_concrete_function().graph.as_graph_def()
+
+    mlir_tf = import_graphdef(
+        tf_graph_def,
+        "tf-standard-pipeline",
+        False,
+        input_names=["lhs", "rhs"],
+        input_data_types=["DT_FLOAT", "DT_FLOAT"],
+        input_data_shapes=["10,10", "10,10"],
+        output_names=["Add"])
+    # Check whether the mlir-function signature has the mentioned
+    # inputs and outputs.
+    self.assertRegex(
+        mlir_tf,
+        r"func @main\(%arg0: tensor<10x10xf32>, %arg1: tensor<10x10xf32>")
+    self.assertRegex(mlir_tf, r'inputs = "lhs,rhs"')
+    self.assertRegex(mlir_tf, r'outputs = "Add"')
+
+    # Same check with scalar input (empty input shape).
+    mlir_tf = import_graphdef(
+        tf_graph_def,
+        "tf-standard-pipeline",
+        False,
+        input_names=["lhs", "rhs"],
+        input_data_types=["DT_FLOAT", "DT_FLOAT"],
+        input_data_shapes=["", ""],
+        output_names=["Add"])
+    self.assertRegex(mlir_tf,
+                     r"func @main\(%arg0: tensor<f32>, %arg1: tensor<f32>")
+
+    # Test invalid test cases where no. of input names is invalid/wrong.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Length of input node array and data type doesn't match"):
+
+      import_graphdef(
+          tf_graph_def,
+          "tf-standard-pipeline",
+          False,
+          input_names=["lhs"],
+          input_data_types=["DT_FLOAT", "DT_FLOAT"],
+          input_data_shapes=["10,10", "10,10"],
+          output_names=["Add"])
+
+    # Test invalid test cases where the input shapes argument is wrong.
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "Dimensions must be equal"):
+
+      import_graphdef(
+          tf_graph_def,
+          "tf-standard-pipeline",
+          False,
+          input_names=["lhs", "rhs"],
+          input_data_types=["DT_FLOAT", "DT_FLOAT"],
+          input_data_shapes=["10,11", "10,10"],
+          output_names=["Add"])
+
+
+class MLIRConcreteFunctionImportTest(test.TestCase):
+
+  @test_util.run_v2_only
+  def testImport(self):
+
+    @def_function.function
+    def sqr(i):
+      return i * i
+
+    concrete_function = sqr.get_concrete_function(
+        tensor_spec.TensorSpec(None, dtypes.float32))
+    mlir_module = mlir.convert_function(concrete_function, show_debug_info=True)
+    self.assertRegex(mlir_module, r'func @.*sqr.*\(')
+    self.assertRegex(mlir_module, r'loc\(".*mlir_test.py":.*:1\)')
+
+  @test_util.run_v2_only
+  def testImportWithCall(self):
+
+    @def_function.function
+    def callee(i):
+      return i
+
+    @def_function.function
+    def caller(i):
+      return callee(i)
+
+    concrete_function = caller.get_concrete_function(
+        tensor_spec.TensorSpec(None, dtypes.float32))
+    mlir_module = mlir.convert_function(concrete_function)
+    self.assertRegex(mlir_module, r'func @.*caller.*\(')
+    self.assertRegex(mlir_module, r'func private @.*callee.*\(')
+
+  @test_util.run_v2_only
+  def testImportWithControlRet(self):
+
+    @def_function.function
+    def logging():
+      logging_ops.print_v2('some message')
+
+    concrete_function = logging.get_concrete_function()
+    mlir_module = mlir.convert_function(concrete_function, pass_pipeline='')
+    self.assertRegex(mlir_module, r'tf\.PrintV2')
+    self.assertRegex(mlir_module, r'tf_executor.fetch.*: !tf_executor.control')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/tensorrt/BUILD

@@ -0,0 +1,173 @@
+# Description:
+#   Wrap NVIDIA TensorRT (http://developer.nvidia.com/tensorrt) with tensorflow
+#   and provide TensorRT operators and converter package.
+#   APIs are meant to change over time.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+# cuda_py_test and cuda_py_tests enable XLA tests by default. We can't
+# combine XLA with TensorRT currently and should set
+# xla_enable_strict_auto_jit to False to disable XLA tests.
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+exports_files(glob([
+    "test/testdata/*",
+]))
+
+py_library(
+    name = "init_py",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    deps = [
+        ":tf_trt_integration_test_base",  # build_cleaner: keep
+        ":trt_convert_py",
+    ],
+)
+
+py_library(
+    name = "trt_convert_py",
+    srcs = ["trt_convert.py"],
+    strict_deps = True,
+    deps = [
+        ":utils",
+        "//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
+        "//tensorflow/compiler/tf2tensorrt:trt_ops_loader",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:wrap_function",
+        "//tensorflow/python/framework:convert_to_constants",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:importer",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/grappler:tf_optimizer",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:resource_variable_ops_gen",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:builder",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/trackable:asset",
+        "//tensorflow/python/trackable:autotrackable",
+        "//tensorflow/python/trackable:resource",
+        "//tensorflow/python/training:saver",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:lazy_loader",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+        "@six_archive//:six",
+    ],
+)
+
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:dtypes",
+        "@pypi//packaging",
+    ],
+)
+
+py_library(
+    name = "tf_trt_integration_test_base",
+    strict_deps = True,
+    deps = [
+        ":trt_convert_py",
+        ":utils",
+        "//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compiler/tensorrt/test:tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:graph_io",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:trace",
+        "//tensorflow/python/saved_model:builder",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/saved_model:utils",
+        "//tensorflow/python/tools:saved_model_utils",
+        "//tensorflow/python/trackable:autotrackable",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "trt_convert_test",
+    srcs = ["trt_convert_test.py"],
+    data = [
+        "//tensorflow/python/compiler/tensorrt/test:trt_convert_test_data",
+    ],
+    tags = [
+        "no_cuda_on_cpu_tap",
+        # TODO(b/297490791): Reenable after TensorRT regression has been fixed
+        "no_oss",
+        "no_pip",
+        "nomac",
+        # TODO(b/303453873): Re-enable tests once TensorRT has been updated
+        "notap",
+    ],
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":trt_convert_py",
+        "//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
+        "//tensorflow/compiler/tf2tensorrt:trt_engine_instance_proto_py",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compiler/tensorrt/test:test_utils",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:convert_to_constants",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:importer",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops_gen",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:builder",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/saved_model:save_options",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/saved_model:utils",
+        "//tensorflow/python/tools:saved_model_utils",
+        "//tensorflow/python/trackable:autotrackable",
+        "//tensorflow/python/util:lazy_loader",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/compiler/tensorrt/README.md

@@ -0,0 +1,60 @@
+# Using TensorRT in TensorFlow (TF-TRT)
+
+Note: Starting from v.2.18.0, TensorFlow doesn't support TensorRT.
+
+This module provides necessary bindings and introduces `TRTEngineOp` operator
+that wraps a subgraph in TensorRT. This module is under active development.
+
+## Installing TF-TRT
+
+Currently TensorFlow nightly builds include TF-TRT by default, which means you
+don't need to install TF-TRT separately. You can pull the latest TF containers
+from docker hub or install the latest TF pip package to get access to the latest
+TF-TRT.
+
+If you want to use TF-TRT on NVIDIA Jetson platform, you can find the download
+links for the relevant TensorFlow pip packages here:
+https://docs.nvidia.com/deeplearning/dgx/index.html#installing-frameworks-for-jetson
+
+## Installing TensorRT
+
+In order to make use of TF-TRT, you will need a local installation of TensorRT.
+Installation instructions for compatibility with TensorFlow are provided on the
+[TensorFlow GPU support](https://www.tensorflow.org/install/gpu) guide.
+
+## Examples
+
+You can find example scripts for running inference on deep learning models in
+this repository: https://github.com/tensorflow/tensorrt
+
+We have used these examples to verify the accuracy and performance of TF-TRT.
+For more information see
+[Verified Models](https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html#verified-models).
+
+## Documentation
+
+[TF-TRT documentation](https://docs.nvidia.com/deeplearning/frameworks/tf-trt-user-guide/index.html)
+gives an overview of the supported functionalities, provides tutorials and
+verified models, explains best practices with troubleshooting guides.
+
+## Tests
+
+TF-TRT includes both Python tests and C++ unit tests. Most of Python tests are
+located in the test directory and they can be executed using `bazel test` or
+directly with the Python command. Most of the C++ unit tests are used to test
+the conversion functions that convert each TF op to a number of TensorRT layers.
+
+## Compilation
+
+In order to compile the module, you need to have a local TensorRT installation
+(libnvinfer.so and respective include files). During the configuration step,
+TensorRT should be enabled and installation path should be set. If installed
+through package managers (deb,rpm), configure script should find the necessary
+components from the system automatically. If installed from tar packages, user
+has to set path to location where the library is installed during configuration.
+
+```shell
+bazel build --config=cuda --config=opt //tensorflow/tools/pip_package:build_pip_package
+bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/
+```
+

tensorflow/python/compiler/tensorrt/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Exposes the python wrapper for TensorRT graph transforms."""
+
+# pylint: disable=unused-import,line-too-long
+from tensorflow.python.compiler.tensorrt import trt_convert as trt
+# pylint: enable=unused-import,line-too-long

tensorflow/python/compiler/tensorrt/model_tests/BUILD

@@ -0,0 +1,73 @@
+# Description:
+#   Run sample models with TensorRT through TF-TRT bridge. Test TensorRT
+#   numerics and latency.
+
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+exports_files(glob([
+    "models/*",
+]))
+
+py_library(
+    name = "model_handler",
+    srcs = ["model_handler.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/compiler/tensorrt:trt_convert_py",
+        "//tensorflow/python/framework:convert_to_constants",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:importer",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "result_analyzer",
+    srcs = ["result_analyzer.py"],
+    strict_deps = True,
+    deps = [
+        ":model_handler",
+        "//tensorflow/python/compiler/tensorrt:trt_convert_py",
+        "//tensorflow/python/platform:tf_logging",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_binary(
+    name = "run_models",
+    srcs = ["run_models.py"],
+    data = ["sample_model/saved_model.pb"],
+    strict_deps = True,
+    deps = [
+        ":model_handler",
+        ":result_analyzer",
+        "//tensorflow/python/compiler/tensorrt:trt_convert_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:tag_constants",
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+    ],
+)

tensorflow/python/compiler/tensorrt/model_tests/model_handler.py

@@ -0,0 +1,670 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Loads, converts, calibrates, and runs sample models."""
+
+import abc
+import collections
+import functools
+import itertools
+import tempfile
+import time
+from typing import Callable, Iterable, List, Mapping, Optional, Sequence, Union
+
+import numpy as np
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.framework import tensor_shape_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python.client import session
+from tensorflow.python.compiler.tensorrt import trt_convert as trt
+from tensorflow.python.framework import convert_to_constants
+from tensorflow.python.framework import dtypes as tf_dtypes
+from tensorflow.python.framework import importer
+from tensorflow.python.framework import ops as framework_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import load as saved_model_load
+from tensorflow.python.saved_model import loader as saved_model_loader
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import tag_constants
+
+
+# pylint: disable=bad-whitespace
+### Helper Functions
+def _remove_graph_sequence_number(name: str) -> str:
+  return name.split(":")[0]
+
+
+def _get_concrete_tensor_shape(
+    tensor_shape: tensor_shape_pb2.TensorShapeProto,
+    batch_size: Optional[int] = None) -> Sequence[int]:
+  """Gets a concrete tensor shape without dynamic dimensions."""
+  if tensor_shape.unknown_rank:
+    raise ValueError("Cannot generates random tensors for unknown rank!")
+  shape = [dim.size for dim in tensor_shape.dim]
+  if not shape:
+    raise ValueError("The tensor cannot have a rank of 0!")
+  if shape[0] < 0:
+    if batch_size is None or batch_size <= 0:
+      raise ValueError("Must provide a valid batch size "
+                       "as the tensor has a dynamic batch size!")
+    shape[0] = batch_size
+  if any(filter(lambda x: x < 0, shape)):
+    raise ValueError("Cannot have dynamic dimensions except for batch size!")
+  return shape
+
+
+def _generate_random_tensor_ops(shape: Sequence[int], dtype: tf_dtypes.DType,
+                                name: str) -> framework_ops.Tensor:
+  # Need to generate a random tensor in float32/int32 and cast to a different
+  # datatype as random_ops doesn't suppprt all the datatypes.
+  random_dtype = tf_dtypes.float32 if dtype.is_floating else tf_dtypes.int32
+  # tf.bool doesn't have `max` attribute
+  dtype_max = 1 if dtype == tf_dtypes.bool else dtype.max
+  return math_ops.cast(
+      random_ops.random_uniform(
+          shape=shape,
+          dtype=random_dtype,
+          # Limits maximum value as 255 to simulate pixel values, avoid
+          # generating large numbers and causing overflows.
+          maxval=min(dtype_max, random_dtype.max, 255),
+      ),
+      dtype=dtype,
+      name=name,
+  )
+
+
+def _generate_random_tensor_v1(tensor_info: meta_graph_pb2.TensorInfo,
+                               batch_size: Optional[int] = None) -> np.ndarray:
+  """Generates a random tensor based on the data type and tensor shape."""
+  dtype = tf_dtypes.as_dtype(tensor_info.dtype)
+  shape = _get_concrete_tensor_shape(tensor_info.tensor_shape, batch_size)
+  with framework_ops.Graph().as_default() as graph, session.Session(
+      graph=graph):
+    return _generate_random_tensor_ops(
+        shape=shape,
+        dtype=dtype,
+        name=_remove_graph_sequence_number(tensor_info.name)).eval()
+
+
+def _generate_random_tensor_v2(
+    tensor: framework_ops.Tensor,
+    batch_size: Optional[int] = None) -> framework_ops.Tensor:
+  """Generates a random tensor based on the data type and tensor shape."""
+  shape = _get_concrete_tensor_shape(tensor.shape.as_proto(), batch_size)
+  return _generate_random_tensor_ops(
+      shape=shape, dtype=tensor.dtype, name=tensor.name)
+
+
+# Models are repeatedly loaded for different TensorRT conversion settings.
+# Using cache can reduce I/O.
+@functools.lru_cache()
+def load_meta_graph(
+    saved_model_dir: str, saved_model_tags: str,
+    saved_model_signature_key: str) -> meta_graph_pb2.MetaGraphDef:
+  """Loads a `tf.MetaGraphDef` in TF1."""
+  with framework_ops.Graph().as_default() as graph, session.Session(
+      graph=graph) as sess:
+    meta_graph = saved_model_loader.load(
+        sess=sess,
+        export_dir=saved_model_dir,
+        tags=saved_model_tags,
+    )
+    output_node_names = [
+        _remove_graph_sequence_number(tensor.name) for tensor in
+        meta_graph.signature_def[saved_model_signature_key].outputs.values()
+    ]
+    graph_def = (
+        convert_to_constants.convert_variables_to_constants_from_session_graph(
+            sess, meta_graph.graph_def, output_node_names))
+    meta_graph.graph_def.CopyFrom(graph_def)
+  return meta_graph
+
+
+@functools.lru_cache()
+def load_graph_func(saved_model_dir: str, saved_model_tags: str,
+                    saved_model_signature_key: str):
+  """Loads a graph function in TF2."""
+  imported = saved_model_load.load(
+      export_dir=saved_model_dir, tags=saved_model_tags)
+  graph_func = imported.signatures[saved_model_signature_key]
+  return convert_to_constants.convert_variables_to_constants_v2(graph_func)
+
+
+### Test Classes
+class ModelConfig(
+    collections.namedtuple("ModelConfig", [
+        "saved_model_dir", "saved_model_tags", "saved_model_signature_key",
+        "default_batch_size"
+    ])):
+  """Configurations for test models."""
+
+  def __new__(cls,
+              saved_model_dir: str,
+              saved_model_tags: Sequence[str] = (tag_constants.SERVING,),
+              saved_model_signature_key: str = (
+                  signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY),
+              default_batch_size: int = 1):
+    return super(ModelConfig,
+                 cls).__new__(cls, saved_model_dir, saved_model_tags,
+                              saved_model_signature_key, default_batch_size)
+
+
+class TestResult(
+    collections.namedtuple("TestResult", [
+        "model_config", "enable_gpu", "output_names", "output_tensors",
+        "model_latency", "trt_convert_params"
+    ])):
+  """Configuration and results for a single model testing."""
+
+  def __new__(cls,
+              model_config: ModelConfig,
+              enable_gpu: bool,
+              output_names: Sequence[str],
+              output_tensors: Sequence[np.ndarray],
+              model_latency: List[float],
+              trt_convert_params: trt.TrtConversionParams = None):
+    return super(TestResult,
+                 cls).__new__(cls, model_config, enable_gpu, output_names,
+                              output_tensors, model_latency, trt_convert_params)
+
+
+class TestResultCollection(
+    collections.namedtuple("TestResultCollection", [
+        "test_name", "model_config", "cpu_base_result", "gpu_base_result",
+        "trt_results"
+    ])):
+  """Configuration and results for a series of model testing."""
+
+  def __new__(cls,
+              test_name: str,
+              model_config: ModelConfig,
+              cpu_base_result: TestResult,
+              gpu_base_result: TestResult,
+              trt_results: Sequence[TestResult] = tuple()):
+    return super(TestResultCollection,
+                 cls).__new__(cls, test_name, model_config, cpu_base_result,
+                              gpu_base_result, trt_results)
+
+  @property
+  def results(self) -> Iterable[TestResult]:
+    return filter(
+        lambda x: x is not None,
+        itertools.chain([self.cpu_base_result, self.gpu_base_result],
+                        self.trt_results))
+
+
+class _ModelHandlerBase(metaclass=abc.ABCMeta):
+  """Base class for running a model."""
+
+  def __init__(self, model_config: ModelConfig):
+    self._model_config = model_config
+
+  def __str__(self) -> str:
+    return str(self._model_config)
+
+  def __repr__(self) -> str:
+    return "{}({})".format(self.__class__.__name__, str(self))
+
+  @property
+  def model_config(self) -> ModelConfig:
+    return self._model_config
+
+  @property
+  def input_tensort_names(self) -> Sequence[str]:
+    """Names of input tensors."""
+
+  @property
+  def output_tensor_names(self) -> Sequence[str]:
+    """Names of output tensors."""
+
+  @abc.abstractmethod
+  def generate_random_inputs(
+      self,
+      batch_size: Optional[int] = None
+  ) -> Mapping[str, Union[np.ndarray, framework_ops.Tensor]]:
+    """Generates mapping from names to input tensors."""
+
+  @abc.abstractmethod
+  def run(self,
+          inputs=None,
+          warmup_iterations: int = 10,
+          benchmark_iterations: int = 100,
+          enable_gpu: bool = True) -> TestResult:
+    """Runs the model with provided or randomly generated input tensors.
+
+    Args:
+      inputs: Mapping from names to input ndarrays in TF1, or a sequence of
+        tensors in TF2. If `None`, ramdomly generated inputs will be used
+        instead.
+      warmup_iterations: Number of inferences to warm up the runtime.
+      benchmark_iterations: Number of inferences to measure the latency.
+      enable_gpu: Whether it is allowed to use GPU or not.
+
+    Returns:
+      `TestResult` summarizing latency and numerics information.
+    """
+
+
+class ModelHandlerV1(_ModelHandlerBase):
+  """Runs a model in TF1."""
+
+  @property
+  def meta_graph(self) -> meta_graph_pb2.MetaGraphDef:
+    return load_meta_graph(
+        saved_model_dir=self.model_config.saved_model_dir,
+        saved_model_tags=self.model_config.saved_model_tags,
+        saved_model_signature_key=self.model_config.saved_model_signature_key)
+
+  @property
+  def input_tensor_info(self) -> Mapping[str, meta_graph_pb2.TensorInfo]:
+    return self.meta_graph.signature_def[
+        self.model_config.saved_model_signature_key].inputs
+
+  @property
+  def output_tensor_info(self) -> Mapping[str, meta_graph_pb2.TensorInfo]:
+    return self.meta_graph.signature_def[
+        self.model_config.saved_model_signature_key].outputs
+
+  @property
+  def input_tensort_names(self) -> Sequence[str]:
+    return [info.name for info in self.input_tensor_info.values()]
+
+  @property
+  def output_tensor_names(self) -> Sequence[str]:
+    return [info.name for info in self.output_tensor_info.values()]
+
+  def generate_random_inputs(self,
+                             batch_size: Optional[int] = None
+                            ) -> Mapping[str, np.ndarray]:
+    batch_size = batch_size or self.model_config.default_batch_size
+    return {
+        tensor_info.name: _generate_random_tensor_v1(tensor_info, batch_size)
+        for tensor_info in self.input_tensor_info.values()
+    }
+
+  def run(self,
+          inputs: Optional[Mapping[str, np.ndarray]] = None,
+          warmup_iterations=10,
+          benchmark_iterations=100,
+          enable_gpu=True) -> TestResult:
+    inputs = inputs or self.generate_random_inputs()
+    config_proto = None
+    if not enable_gpu:
+      config_proto = config_pb2.ConfigProto(device_count={"CPU": 1, "GPU": 0})
+    logging.info("Running model inference!")
+    with framework_ops.Graph().as_default():
+      with session.Session(config=config_proto) as sess:
+        importer.import_graph_def(self.meta_graph.graph_def, name="")
+        try:
+          output_tensor_names = self.output_tensor_names
+          for _ in range(warmup_iterations):
+            sess.run(fetches=output_tensor_names, feed_dict=inputs)
+          latency = []
+          for _ in range(benchmark_iterations):
+            before = time.time()
+            outputs = sess.run(fetches=output_tensor_names, feed_dict=inputs)
+            latency.append(time.time() - before)
+        except Exception as exc:
+          raise RuntimeError("Failed to run model inference! "
+                             "Model information: {}".format(str(self))) from exc
+    return TestResult(
+        model_config=self.model_config,
+        enable_gpu=enable_gpu,
+        model_latency=latency,
+        output_names=self.output_tensor_names,
+        output_tensors=outputs)
+
+
+class ModelHandlerV2(_ModelHandlerBase):
+  """Runs a model in TF2."""
+
+  @property
+  def graph_func(self):
+    try:
+      return self._graph_func
+    except:
+      graph_func = load_graph_func(
+          saved_model_dir=self.model_config.saved_model_dir,
+          saved_model_tags=self.model_config.saved_model_tags,
+          saved_model_signature_key=self.model_config.saved_model_signature_key)
+      self._graph_func = convert_to_constants.convert_variables_to_constants_v2(
+          graph_func)
+      return self._graph_func
+
+  @property
+  def input_tensor_names(self):
+    return [tensor.name for tensor in self.graph_func.inputs]
+
+  @property
+  def output_tensor_names(self):
+    return [tensor.name for tensor in self.graph_func.outputs]
+
+  def generate_random_inputs(self,
+                             batch_size: Optional[int] = None
+                            ) -> Sequence[framework_ops.Tensor]:
+    batch_size = batch_size or self.model_config.default_batch_size
+    return [
+        _generate_random_tensor_v2(tensor, batch_size)
+        for tensor in self.graph_func.inputs
+    ]
+
+  def run(self,
+          inputs: Optional[Sequence[framework_ops.Tensor]] = None,
+          warmup_iterations=10,
+          benchmark_iterations=100,
+          enable_gpu=True) -> TestResult:
+    inputs = inputs or self.generate_random_inputs()
+    try:
+      device = "/device:gpu:0" if enable_gpu else "/device:cpu:0"
+      with framework_ops.device(device):
+        for _ in range(warmup_iterations):
+          self.graph_func(*inputs)
+        latency = []
+        for _ in range(benchmark_iterations):
+          before = time.time()
+          outputs = self.graph_func(*inputs)
+          latency.append(time.time() - before)
+    except Exception as exc:
+      raise RuntimeError("Failed to run model inference! "
+                         "Model information: {}".format(str(self))) from exc
+    return TestResult(
+        model_config=self.model_config,
+        enable_gpu=enable_gpu,
+        model_latency=latency,
+        output_names=self.output_tensor_names,
+        output_tensors=outputs)
+
+
+class _TrtModelHandlerBase(_ModelHandlerBase):
+  """Base class for converting and running a model."""
+
+  def __init__(
+      self,
+      model_config: ModelConfig,
+      trt_convert_params: trt.TrtConversionParams,
+  ):
+    super(_TrtModelHandlerBase, self).__init__(model_config)
+    self._trt_convert_params = trt_convert_params
+
+    self._converter = self._create_converter(trt_convert_params)
+    self._conversion_is_saved = False
+
+  @abc.abstractmethod
+  def _create_converter(self, trt_convert_params: trt.TrtConversionParams):
+    """Creates a converter for the corresponding TF version."""
+
+  @abc.abstractmethod
+  def _check_conversion(self, conversion_output):
+    """Checks if conversion output has any TensorRT engines."""
+
+  def _check_contains_trt_engine(self, graph_def: graph_pb2.GraphDef):
+    if "TRTEngineOp" not in [node.op for node in graph_def.node]:
+      raise RuntimeError("Failed to convert to TensorRT! "
+                         "Model Information: {}".format(str(self)))
+
+  def __str__(self) -> str:
+    base = super(_TrtModelHandlerBase, self).__str__()
+    return "{}, TrtConversionParams: {}".format(base,
+                                                str(self._trt_convert_params))
+
+  @property
+  def trt_convert_params(self) -> trt.TrtConversionParams:
+    return self._trt_convert_params
+
+  @abc.abstractmethod
+  def convert(self,
+              calibration_inputs: Optional[Mapping[str, np.ndarray]] = None,
+              num_runs=1) -> None:
+    """Converts the model with TensorRT and calibrates if using INT8 precision mode.
+
+    Args:
+      calibration_inputs: Mapping from input names to ndarrays in TF1. Or a
+        sequence of tensors in TF2. Used as calibration data.
+      num_runs: Number of calibration runs.
+    """
+
+  def save(self,
+           output_saved_model_dir: Optional[str] = None,
+           overwrite=True) -> None:
+    """Saves a TensorRT converted model."""
+    if self._conversion_is_saved and not overwrite:
+      return
+    output_saved_model_dir = output_saved_model_dir or tempfile.mkdtemp()
+    logging.info("Saving TensorRT model to %s!", output_saved_model_dir)
+    self._converter.save(output_saved_model_dir)
+    self._model_config = self.model_config._replace(
+        saved_model_dir=output_saved_model_dir)
+    self._conversion_is_saved = True
+
+
+class TrtModelHandlerV1(_TrtModelHandlerBase, ModelHandlerV1):
+  """Converts a TF1 model with TensorRT and runs the converted model."""
+
+  def _create_converter(self, trt_convert_params: trt.TrtConversionParams):
+    conversion_nodes_denylist = self.output_tensor_names
+    return trt.TrtGraphConverter(
+        input_saved_model_dir=self.model_config.saved_model_dir,
+        input_saved_model_tags=self.model_config.saved_model_tags,
+        input_saved_model_signature_key=(
+            self.model_config.saved_model_signature_key),
+        nodes_denylist=conversion_nodes_denylist,
+        max_workspace_size_bytes=trt_convert_params.max_workspace_size_bytes,
+        precision_mode=trt_convert_params.precision_mode,
+        minimum_segment_size=trt_convert_params.minimum_segment_size,
+        maximum_cached_engines=trt_convert_params.maximum_cached_engines,
+        use_calibration=trt_convert_params.use_calibration,
+        max_batch_size=self.model_config.default_batch_size,
+        is_dynamic_op=False,
+    )
+
+  _check_conversion = _TrtModelHandlerBase._check_contains_trt_engine
+
+  def convert(self,
+              calibration_inputs: Optional[Mapping[str, np.ndarray]] = None,
+              num_runs=1) -> None:
+    logging.info("Converting with TensorRT!")
+    self._check_conversion(self._converter.convert())
+
+    if (self.trt_convert_params.precision_mode == trt.TrtPrecisionMode.INT8 and
+        self.trt_convert_params.use_calibration):
+      logging.info("Calibrating with TensorRT!")
+      if not calibration_inputs:
+        raise ValueError("Must provide calibration data "
+                         "when using TensorRT calibration!")
+      try:
+        self._converter.calibrate(
+            fetch_names=self.output_tensor_names,
+            num_runs=num_runs,
+            feed_dict_fn=lambda: calibration_inputs)
+      except Exception as exc:
+        raise RuntimeError("Failed to calibrate! "
+                           "Model Information: {}".format(str(self))) from exc
+
+  def run(self,
+          inputs: Optional[Mapping[str, np.ndarray]] = None,
+          warmup_iterations=10,
+          benchmark_iterations=100) -> TestResult:
+    self.save(overwrite=False)
+    self._check_conversion(self.meta_graph.graph_def)
+    logging.info("Running with TensorRT!")
+    test_result = ModelHandlerV1.run(
+        self, inputs, warmup_iterations, benchmark_iterations, enable_gpu=True)
+    return test_result._replace(trt_convert_params=self._trt_convert_params)
+
+
+class TrtModelHandlerV2(_TrtModelHandlerBase, ModelHandlerV2):
+  """Converts a TF2 model with TensorRT and runs the converted model."""
+
+  def _create_converter(self, trt_convert_params: trt.TrtConversionParams):
+    return trt.TrtGraphConverterV2(
+        input_saved_model_dir=self.model_config.saved_model_dir,
+        input_saved_model_tags=self.model_config.saved_model_tags,
+        input_saved_model_signature_key=(
+            self.model_config.saved_model_signature_key),
+        **trt_convert_params._asdict())
+
+  def _check_conversion(self, graph_func):
+    graph_def = graph_func.graph.as_graph_def()
+    self._check_contains_trt_engine(graph_def)
+
+  def convert(self,
+              calibration_inputs: Optional[Sequence[
+                  framework_ops.Tensor]] = None,
+              num_runs=1) -> None:
+    logging.info("Converting with TensorRT!")
+
+    calibration_input_fn = None
+    if (self.trt_convert_params.precision_mode == trt.TrtPrecisionMode.INT8 and
+        self.trt_convert_params.use_calibration):
+      logging.info("Calibrating with TensorRT at the same time!")
+      if not calibration_inputs:
+        raise ValueError("Must provide calibration data "
+                         "when using TensorRT calibration!")
+
+      def gets_calibration_input():
+        for _ in range(num_runs):
+          yield calibration_inputs
+
+      calibration_input_fn = gets_calibration_input
+
+    self._check_conversion(self._converter.convert(calibration_input_fn))
+
+  def run(self,
+          inputs: Optional[Sequence[framework_ops.Tensor]] = None,
+          warmup_iterations=10,
+          benchmark_iterations=100) -> TestResult:
+    self.save(overwrite=False)
+    self._check_conversion(self.graph_func)
+    logging.info("Running with TensorRT!")
+    test_result = ModelHandlerV2.run(
+        self, inputs, warmup_iterations, benchmark_iterations, enable_gpu=True)
+    return test_result._replace(trt_convert_params=self._trt_convert_params)
+
+
+class _ModelHandlerManagerBase(metaclass=abc.ABCMeta):
+  """Manages a series of ModelHandlers for aggregated testing/benchmarking."""
+
+  def __init__(
+      self, name: str, model_config: ModelConfig,
+      default_trt_convert_params: trt.TrtConversionParams,
+      trt_convert_params_updater: Callable[[trt.TrtConversionParams],
+                                           Iterable[trt.TrtConversionParams]]):
+    self._ori_model = self.model_handler_cls(model_config)
+    self._trt_models = []
+    for trt_convert_params in trt_convert_params_updater(
+        default_trt_convert_params):
+      trt_model = self.trt_model_handler_cls(
+          model_config, trt_convert_params=trt_convert_params)
+      self._trt_models.append(trt_model)
+
+    self._name = name
+    self._result_collection = None
+
+  def __str__(self) -> str:
+    return "Input Model: {}".format(str(self._ori_model))
+
+  def __repr__(self) -> str:
+    return "{}({})".format(self.__class__.__name__, str(self))
+
+  @property
+  @classmethod
+  @abc.abstractmethod
+  def model_handler_cls(cls):
+    """The model handler class. ModelHandleV1/ModelHandlerV2."""
+
+  @property
+  @classmethod
+  @abc.abstractmethod
+  def trt_model_handler_cls(cls):
+    """The TensorRT model handler class. TrtModelHandleV1/TrtModelHandlerV2."""
+
+  @property
+  def name(self) -> str:
+    return self._name
+
+  @property
+  def model_config(self) -> ModelConfig:
+    return self._ori_model.model_config
+
+  def generate_random_inputs(self, batch_size: Optional[int] = None):
+    return self._ori_model.generate_random_inputs(batch_size)
+
+  def convert(self, calibration_inputs=None, num_runs=1) -> None:
+    """Converts models with TensorRT and calibrates if using INT8 precision mode.
+
+    Args:
+      calibration_inputs: Mapping from input names to ndarrays in TF1. Or a
+        sequence of tensors in TF2. Used as calibration data.
+      num_runs: Number of calibration runs.
+    """
+    for trt_model in self._trt_models:
+      trt_model.convert(calibration_inputs, num_runs)
+
+  def run(self,
+          inputs=None,
+          warmup_iterations: int = 10,
+          benchmark_iterations: int = 100) -> TestResultCollection:
+    """Runs model inference with provided or randomly generated input tensors.
+
+    Args:
+      inputs: Mapping from names to input ndarrays in TF1. Or a sequence of
+        tensors in TF2. If `None`, ramdomly generated input tensors will be used
+        instead.
+      warmup_iterations: Number of inferences to warm up the runtime.
+      benchmark_iterations: Number of inferences to measure the latency.
+
+    Returns:
+      `TestResultCollection` summarizing latency and numerics information for
+      different TensorRT conversion settings.
+    """
+    inputs = inputs or self.generate_random_inputs()
+
+    def run_model(model, **kwargs):
+      return model.run(inputs, warmup_iterations, benchmark_iterations,
+                       **kwargs)
+
+    # Some models include operations that can only run on GPU.
+    try:
+      cpu_base_result = run_model(self._ori_model, enable_gpu=False)
+    except RuntimeError as err:
+      logging.info("%s cannot run on CPU. Reason: %s.",
+                   self._ori_model.model_config, err)
+      cpu_base_result = None
+    gpu_base_result = run_model(self._ori_model, enable_gpu=True)
+    trt_results = list(map(run_model, self._trt_models))
+
+    return TestResultCollection(
+        test_name=self._name,
+        model_config=self.model_config,
+        cpu_base_result=cpu_base_result,
+        gpu_base_result=gpu_base_result,
+        trt_results=trt_results)
+
+
+class ModelHandlerManagerV1(_ModelHandlerManagerBase):
+  """Manages a series of ModelHandlers for aggregated testing/benchmarking in TF1."""
+
+  model_handler_cls = ModelHandlerV1
+  trt_model_handler_cls = TrtModelHandlerV1
+
+
+class ModelHandlerManagerV2(_ModelHandlerManagerBase):
+  """Manages a series of ModelHandlers for aggregated testing/benchmarking in TF2."""
+
+  model_handler_cls = ModelHandlerV2
+  trt_model_handler_cls = TrtModelHandlerV2

tensorflow/python/compiler/tensorrt/model_tests/result_analyzer.py

@@ -0,0 +1,261 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Analyzes the latency and numerics information of sample model inference."""
+
+import itertools
+import json
+from typing import Any, Callable, Optional, Sequence, Tuple, Union
+from collections import namedtuple
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.model_tests import model_handler
+import tensorflow.python.compiler.tensorrt.trt_convert as trt
+from tensorflow.python.platform import tf_logging as logging
+
+
+# pylint: disable=bad-whitespace
+class DataFrame:
+  """Lightweight immutable Dataframe similar to Pandas Dataframe."""
+
+  def __init__(self,
+               column_names: Sequence[str],
+               rows: Sequence[Sequence[Any]] = None,
+               columns: Sequence[Sequence[Any]] = None):
+    self._column_names = column_names
+    if not rows and not columns:
+      raise ValueError("Cannot initialize with empty data!")
+    self._rows = rows
+    self._columns = columns
+
+  @property
+  def n_rows(self) -> int:
+    return len(self._rows) if self._rows else len(self._columns[0])
+
+  @property
+  def n_columns(self) -> int:
+    return len(self._columns) if self._columns else len(self._rows[0])
+
+  @property
+  def column_names(self) -> Sequence[str]:
+    return self._column_names
+
+  @property
+  def rows(self) -> Sequence[Sequence[Any]]:
+    return self._rows if self._rows else [
+        [c[i] for c in self._columns] for i in range(len(self._columns[0]))
+    ]
+
+  @property
+  def columns(self) -> Sequence[Sequence[Any]]:
+    return self._columns if self._columns else [
+        [r[i] for r in self._rows] for i in range(len(self._rows[0]))
+    ]
+
+  def __add__(self, other: "DataFrame") -> "DataFrame":
+    if (not set(self.column_names).intersection(other.column_names) and
+        len(self.rows) == len(other.rows)):
+      return DataFrame(
+          column_names=list(
+              itertools.chain(self.column_names, other.column_names)),
+          columns=list(itertools.chain(self.columns, other.columns)))
+    if self.column_names == other.column_names:
+      return DataFrame(
+          column_names=self.column_names,
+          rows=list(itertools.chain(self.rows, other.rows)))
+    raise ValueError("Cannot combine two DataFrame")
+
+  def __iadd__(self, other: "DataFrame") -> "DataFrame":
+    tmp = self + other
+    self._column_names = tmp._column_names
+    self._rows, self._columns = tmp._rows, tmp._columns
+    return self
+
+  def __call__(self, r: int, c: Optional[Union[int, str]] = None) -> Any:
+    if c is None:
+      return dict(zip(self.column_names, self.rows[r]))
+    c = self._column_names.index(c) if isinstance(c, str) else c
+    return self._rows[r][c] if self._rows else self._columns[c][r]
+
+  def __str__(self) -> str:
+    return ",".join(self.column_names) + "\n" + "\n".join(",".join(
+        "N/A" if v is None else str(v) for v in row) for row in self.rows)
+
+  def to_csv(self, path: str):
+    with open(path, "w") as file:
+      file.write(str(self))
+
+  def to_json(self, path: str):
+    with open(path, "w") as file:
+      json.dump([dict(zip(self.column_names, r)) for r in self.rows], file)
+
+
+def extract_test_info(
+    test_results: model_handler.TestResultCollection) -> DataFrame:
+  """Extracts the test information."""
+  column_names = list(
+      itertools.chain(model_handler.ModelConfig._fields,
+                      ["enable_gpu", "trt_model"],
+                      trt.TrtConversionParams._fields))
+  rows = []
+  for result in test_results.results:
+    r = list(result.model_config) + [result.enable_gpu]
+    if result.trt_convert_params is not None:
+      r += [True] + list(result.trt_convert_params)
+    else:
+      r += [False] + [None for _ in trt.TrtConversionParams._fields]
+    rows.append(r)
+  return DataFrame(column_names=column_names, rows=rows)
+
+
+def analyze_test_latency(test_results: model_handler.TestResultCollection,
+                         use_cpu_baseline: bool) -> DataFrame:
+  """Analyzes test latency."""
+  base_result = (
+      test_results.cpu_base_result
+      if use_cpu_baseline else test_results.gpu_base_result)
+  if base_result is None:
+    raise ValueError(
+        f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!")
+  base_mean_time = np.mean(base_result.model_latency).item()
+  column_names = ["time(ms)", "speedup"]
+  rows = []
+  for result in test_results.results:
+    mean_time = np.mean(result.model_latency).item()
+    rows.append([mean_time * 1000.0, base_mean_time / mean_time])
+  return DataFrame(column_names=column_names, rows=rows)
+
+
+def str_histogram(buffer, desc):
+  hist, bin_edges = np.histogram(buffer)
+  max_num_elems = np.amax(hist)
+  bin_edges = ["{:.3g}".format(bin) for bin in bin_edges]
+  max_start_bin_width = max(len(bin) for bin in bin_edges)
+  max_end_bin_width = max(len(bin) for bin in bin_edges[1:])
+  MAX_WIDTH = 40
+  ret = "\n========================================================\n"
+  ret += "**** Output " + desc + " ****\n"
+  ret += "---- Histogram ----\n"
+  ret += "{:{width}}|  Num Elems | Visualization\n".format(
+      "Bin Range", width=max_start_bin_width + max_end_bin_width + 5)
+  for num, bin_start, bin_end in zip(hist, bin_edges, bin_edges[1:]):
+    bar = "#" * int(MAX_WIDTH * float(num) / float(max_num_elems))
+    ret += ("({:<{max_start_bin_width}}, {:<{max_end_bin_width}}) | {:10} | "
+            "{:}\n").format(
+        bin_start,
+        bin_end,
+        num,
+        bar,
+        max_start_bin_width=max_start_bin_width,
+        max_end_bin_width=max_end_bin_width,
+    )
+  return ret
+
+
+def analyze_test_numerics(test_results: model_handler.TestResultCollection,
+                          use_cpu_baseline: bool) -> (DataFrame, str):
+  """Analyzes test numerics."""
+  preprocess_funcs = {"abs_diff": lambda x, y: np.fabs(x - y)}
+  postprocess_funcs = {"mean": np.mean}
+  column_names = []
+  columns = []
+  base_result = (
+      test_results.cpu_base_result
+      if use_cpu_baseline else test_results.gpu_base_result)
+  if base_result is None:
+    raise ValueError(
+        f"No {'CPU' if use_cpu_baseline else 'GPU'} baseline found!")
+  for fn0, fn1 in itertools.product(preprocess_funcs, postprocess_funcs):
+    func0, func1 = preprocess_funcs[fn0], postprocess_funcs[fn1]
+    column_names.append("{}_{}".format(fn0, fn1))
+    columns.append([])
+    for result in test_results.results:
+      columns[-1].append(dict())
+      for idx, tensor in enumerate(result.output_tensors):
+        name = base_result.output_names[idx]
+        cpu_tensor = base_result.output_tensors[idx]
+        absdiff = func0(tensor, cpu_tensor)
+        metric_value = func1(absdiff).item()
+        cpu_tensor_hist = str_histogram(cpu_tensor, "cpu_tensor")
+        gpu_tensor_hist = str_histogram(tensor, "gpu_tensor")
+        abs_diff_hist = str_histogram(absdiff, "abs_diff")
+        hist_data = (cpu_tensor_hist, gpu_tensor_hist, abs_diff_hist)
+        columns[-1][-1][name] = metric_value
+  return DataFrame(column_names=column_names, columns=columns), hist_data
+
+
+def check_column(df: DataFrame, row: int, name: str,
+                 fn: Callable[[float], bool]) -> bool:
+  """Checks the values of a column using a custom function and logs abnormals.
+
+  The check is only performed on TensorRT models, not native CPU/GPU models.
+
+  Args:
+    df: The DataFrame to be checked.
+    row: The row in the DataFrame
+    name: The name of the column to be checked.
+    fn: The function that takes a value of at the specified column and returns
+      if the value satisfies the check.
+
+  Returns:
+    Whether all the values of the specified column satisfies the provided check.
+  """
+  is_ok = True
+  if df(row, "trt_model"):
+    if not fn(df(row, name)):
+      logging.error("Unsatisfied %s found at: %s", name, df(row))
+      is_ok = False
+  return is_ok
+
+
+class ResultAnalyzer:
+  """Analyzes ModelHandlerManager results."""
+
+  def __init__(
+      self,
+      use_cpu_latency_baseline: bool,
+      use_cpu_numerics_baseline: bool,
+      perf_checkers: Sequence[Callable[[DataFrame], bool]],
+      acc_checkers: Sequence[Callable[[DataFrame], bool]],
+  ):
+    self._use_cpu_latency_baseline = use_cpu_latency_baseline
+    self._use_cpu_numerics_baseline = use_cpu_numerics_baseline
+    self._perf_checkers = perf_checkers
+    self._acc_checkers = acc_checkers
+
+
+  def analysis(
+      self, test_results: model_handler.TestResultCollection
+  ) -> Tuple[DataFrame, Sequence[bool]]:
+    df = extract_test_info(test_results)
+    df += analyze_test_latency(test_results, self._use_cpu_latency_baseline)
+    df_acc, acc_hist = analyze_test_numerics(test_results,
+                                             self._use_cpu_numerics_baseline)
+    df += df_acc
+    # Index 0 and 1 are non trt_model results, so ignore them here.
+    df_analysis_config = namedtuple("df_analysis_config",
+                                    "precision df_row_index")
+    checker_config = [
+        df_analysis_config("FP32", df.n_rows - 3),
+        df_analysis_config("FP16", df.n_rows - 2),
+        df_analysis_config("INT8", df.n_rows - 1)
+    ]
+
+    checks = []
+    for cc in checker_config:
+      checks.append(self._perf_checkers[cc.precision](df, cc.df_row_index))
+      checks.append(self._acc_checkers[cc.precision](df, cc.df_row_index))
+
+    return df, checks, acc_hist

tensorflow/python/compiler/tensorrt/model_tests/run_models.py

@@ -0,0 +1,287 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Runs sample models with TensorRT and analyzes latency and numerics information."""
+
+import functools
+import os
+import tempfile
+from typing import Callable, Iterable, Sequence
+from collections import namedtuple
+
+from absl import app
+from absl import flags
+
+from tensorflow.python.compiler.tensorrt import trt_convert as trt
+from tensorflow.python.compiler.tensorrt.model_tests import model_handler
+from tensorflow.python.compiler.tensorrt.model_tests import result_analyzer
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config as framework_config
+from tensorflow.python.framework import ops as framework_ops
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test as platform_test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import tag_constants
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    "saved_model_dir",
+    platform_test.test_src_dir_path(
+        "python/compiler/tensorrt/model_tests/sample_model"),
+    "The directory to the testing SavedModel.")
+
+flags.DEFINE_string("saved_model_signature_key",
+                    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
+                    "The signature key of the testing SavedModel being used.")
+
+flags.DEFINE_multi_string("saved_model_tags", (tag_constants.SERVING,),
+                          "The tags of the testing SavedModel being used.")
+
+flags.DEFINE_integer("batch_size", 128,
+                     "The batch size used to run the testing model with.")
+
+flags.DEFINE_boolean("use_tf2", True,
+                     "Whether to test with TF2 behavior or not (TF1).")
+
+flags.DEFINE_boolean("use_int8", True,
+                     "Whether to convert with INT8 precision.")
+
+flags.DEFINE_enum("latency_baseline", "GPU", ["CPU", "GPU"],
+                  "The baseline version for latency improvement analysis.")
+
+flags.DEFINE_enum("numerics_baseline", "CPU", ["CPU", "GPU"],
+                  "The baseline version for numerical difference analysis.")
+
+flags.DEFINE_float(
+    "fp32_speedup_tolerance", 0.90,
+    "Log errors whenever mean TensorRT fp32 speedup is lower than the tolerance."
+)
+
+flags.DEFINE_float(
+    "fp16_speedup_tolerance", 0.90,
+    "Log errors whenever mean TensorRT fp16 speedup is lower than the tolerance."
+)
+
+flags.DEFINE_float(
+    "int8_speedup_tolerance", 0.90,
+    "Log errors whenever mean TensorRT int8 speedup is lower than the tolerance."
+)
+
+flags.DEFINE_float(
+    "fp32_abs_tolerance", 1e-05,
+    "Log errors whenever mean TensorRT fp32 absolute difference is larger than "
+    "the tolerance.")
+
+flags.DEFINE_float(
+    "fp16_abs_tolerance", 5e-02,
+    "Log errors whenever mean TensorRT fp16 absolute difference is larger than "
+    "the tolerance.")
+
+flags.DEFINE_float(
+    "int8_abs_tolerance", 5e-01,
+    "Log errors whenever mean TensorRT int8 absolute difference is larger than "
+    "the tolerance.")
+
+flags.DEFINE_integer(
+    "gpu_memory_limit_mb", None,
+    "Limitation on the device memory being used during TensorRT compilation "
+    "and inference.")
+
+flags.DEFINE_string("output_dir", None, "Output directory of analysis results.")
+
+flags.DEFINE_enum("output_format", "CSV", ["CSV", "JSON"],
+                  "Output format of analysis results.")
+
+DEFAUL_TRT_CONVERT_PARAMS = trt.DEFAULT_TRT_CONVERSION_PARAMS
+
+
+# pylint: disable=bad-whitespace
+def set_up_gpu_memory_limit(memory_limit_mb: int) -> None:
+  gpus = framework_config.list_physical_devices("GPU")
+  virtual_device_config = context.LogicalDeviceConfiguration(
+      memory_limit=memory_limit_mb)
+  for gpu in gpus:
+    framework_config.set_logical_device_configuration(gpu,
+                                                      [virtual_device_config])
+
+
+class SampleRunner(object):
+  """The driver to run all sample models in all specified configurations."""
+
+  def __init__(self, saved_model_dir: str, saved_model_tags: Sequence[str],
+               saved_model_signature_key: str, batch_size: int, output_dir: str,
+               output_format: str, use_tf2: bool, use_int8: bool,
+               analyzer: result_analyzer.ResultAnalyzer):
+    self._output_dir = output_dir or tempfile.mkdtemp(
+        prefix="tf2trt_model_tests")
+    logging.info("Use output directory as: %s", self._output_dir)
+    self._output_format = output_format
+    # The model_configs contains (saved_model_dir, saved_model_signature_key,
+    # batch_size) for each model
+    self._configs = (model_handler.ModelConfig(
+        saved_model_dir=saved_model_dir,
+        saved_model_tags=tuple(saved_model_tags),
+        saved_model_signature_key=saved_model_signature_key,
+        default_batch_size=batch_size),)
+    self._model_handler_manager_cls = (
+        model_handler.ModelHandlerManagerV2
+        if use_tf2 else model_handler.ModelHandlerManagerV1)
+
+    if use_int8:
+      self._precision_modes = [
+          trt.TrtPrecisionMode.FP32, trt.TrtPrecisionMode.FP16,
+          trt.TrtPrecisionMode.INT8]
+    else:
+      self._precision_modes = [
+          trt.TrtPrecisionMode.FP32, trt.TrtPrecisionMode.FP16]
+
+    self._analyzer = analyzer
+
+  def _write_analysis_result(self, df: result_analyzer.DataFrame,
+                             path: str) -> None:
+    if self._output_format == "CSV":
+      df.to_csv(os.path.join(path, "result.csv"))
+    elif self._output_format == "JSON":
+      df.to_json(os.path.join(path, "result.json"))
+    else:
+      raise NotImplementedError("Unsupported output format: {}".format(
+          self._output_format))
+
+  def _run_impl(
+      self, test_name: str,
+      default_trt_converter_params: trt.TrtConversionParams,
+      trt_converter_params_updater: Callable[[trt.TrtConversionParams],
+                                             Iterable[trt.TrtConversionParams]]
+  ) -> None:
+    """Runs all sample models based on a key varying parameter."""
+    for model_config in self._configs:
+      # Loads, compiles, calibrates and runs models.
+      manager = self._model_handler_manager_cls(
+          name=test_name,
+          model_config=model_config,
+          default_trt_convert_params=default_trt_converter_params,
+          trt_convert_params_updater=trt_converter_params_updater)
+      inputs = manager.generate_random_inputs()
+      # As all the data are randomly generated, directly use inference data as
+      # calibration data to produce reliable dynamic ranges.
+      manager.convert(inputs)
+      test_results = manager.run(inputs)
+
+      # Analyzes the latency and numerical results.
+      analysis_result_df, _, acc_hist = self._analyzer.analysis(test_results)
+
+      # Outputs the analysis results
+      model_name = os.path.split(manager.model_config.saved_model_dir)[-1]
+      model_dir = os.path.join(self._output_dir, model_name)
+      gfile.MkDir(model_dir)
+      test_dir = os.path.join(model_dir, test_name)
+      gfile.MkDir(test_dir)
+      with gfile.Open(
+          os.path.join(test_dir, "default_tensorrt_params.txt"), "w") as f:
+        f.write(repr(default_trt_converter_params))
+      with gfile.Open(os.path.join(test_dir, "accuracy_histograms.txt"),
+                      "w") as f:
+        [f.write(h) for h in acc_hist]
+      self._write_analysis_result(analysis_result_df, test_dir)
+
+  def run_trt_precision_tests(self) -> None:
+    """Runs tests for all TensorRT precisions."""
+
+    def trt_converter_params_updater(params: trt.TrtConversionParams):
+      for precision_mode in self._precision_modes:
+        yield params._replace(
+            precision_mode=precision_mode,
+            use_calibration=(precision_mode == trt.TrtPrecisionMode.INT8))
+
+    self._run_impl(
+        test_name="precision_mode_test",
+        default_trt_converter_params=DEFAUL_TRT_CONVERT_PARAMS,
+        trt_converter_params_updater=trt_converter_params_updater)
+
+  def run_all_tests(self) -> None:
+    """Runs all tests available."""
+    self.run_trt_precision_tests()
+    logging.info("Check analysis result at: %s", self._output_dir)
+
+
+def main(argv):
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "False"
+
+  if FLAGS.use_tf2:
+    logging.info("Running in TF2 mode. Eager execution is enabled.")
+    framework_ops.enable_eager_execution()
+  else:
+    logging.info("Running in TF1 mode. Eager execution is disabled.")
+    framework_ops.disable_eager_execution()
+
+  if FLAGS.use_int8:
+    logging.info("Will try converting with INT8 precision.")
+  else:
+    logging.info("Will not try converting with INT8 precision.")
+
+  if FLAGS.gpu_memory_limit_mb:
+    set_up_gpu_memory_limit(FLAGS.gpu_memory_limit_mb)
+
+  tol = namedtuple("tol", "perf acc")
+  tolerances = {
+      trt.TrtPrecisionMode.FP32:
+          tol(perf=float(FLAGS.fp32_speedup_tolerance),
+              acc=float(FLAGS.fp32_abs_tolerance)),
+      trt.TrtPrecisionMode.FP16:
+          tol(perf=float(FLAGS.fp16_speedup_tolerance),
+              acc=float(FLAGS.fp16_abs_tolerance)),
+      trt.TrtPrecisionMode.INT8:
+          tol(perf=float(FLAGS.int8_speedup_tolerance),
+              acc=float(FLAGS.int8_abs_tolerance)),
+  }
+
+  analyzer = result_analyzer.ResultAnalyzer(
+      use_cpu_latency_baseline=FLAGS.latency_baseline == "CPU",
+      use_cpu_numerics_baseline=FLAGS.numerics_baseline == "CPU",
+      perf_checkers={
+          precision: functools.partial(
+              result_analyzer.check_column,
+              name="speedup",
+              fn=lambda x: x > tol.perf)
+          for precision, tol in tolerances.items()
+      },
+      acc_checkers={
+          precision: functools.partial(
+              result_analyzer.check_column,
+              name="abs_diff_mean",
+              fn=lambda x: all(v < tol.acc for v in x.values()))
+          for precision, tol in tolerances.items()
+      })
+
+  runner = SampleRunner(
+      saved_model_dir=FLAGS.saved_model_dir,
+      saved_model_tags=FLAGS.saved_model_tags,
+      saved_model_signature_key=FLAGS.saved_model_signature_key,
+      batch_size=FLAGS.batch_size,
+      output_dir=FLAGS.output_dir,
+      output_format=FLAGS.output_format,
+      use_tf2=FLAGS.use_tf2,
+      use_int8=FLAGS.use_int8,
+      analyzer=analyzer)
+
+  runner.run_all_tests()
+
+
+if __name__ == "__main__":
+  app.run(main)

tensorflow/python/compiler/tensorrt/test/BUILD

@@ -0,0 +1,587 @@
+# Description:
+#   Wrap NVIDIA TensorRT (http://developer.nvidia.com/tensorrt) with tensorflow
+#   and provide TensorRT operators and converter package.
+#   APIs are meant to change over time.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+exports_files(
+    glob(["*test.py"]),
+)
+
+py_library(
+    name = "test_utils",
+    srcs = ["test_utils.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "tf_trt_integration_test_base_srcs",
+    srcs = ["tf_trt_integration_test_base.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow/python/compiler/tensorrt:__pkg__"],
+    deps = [
+        "//tensorflow/compiler/tf2tensorrt:_pywrap_py_utils",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compiler/tensorrt:trt_convert_py",
+        "//tensorflow/python/compiler/tensorrt:utils",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:graph_io",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:trace",
+        "//tensorflow/python/saved_model:builder",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/saved_model:utils",
+        "//tensorflow/python/tools:saved_model_utils",
+        "//tensorflow/python/trackable:autotrackable",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+    ],
+)
+
+filegroup(
+    name = "trt_convert_test_data",
+    srcs = [
+        "testdata/tf_readvariableop_saved_model/saved_model.pb",
+        "testdata/tf_readvariableop_saved_model/variables/variables.data-00000-of-00001",
+        "testdata/tf_readvariableop_saved_model/variables/variables.index",
+        "testdata/tf_variablev2_saved_model/saved_model.pb",
+        "testdata/tf_variablev2_saved_model/variables/variables.data-00000-of-00001",
+        "testdata/tf_variablev2_saved_model/variables/variables.index",
+        "testdata/tftrt_2.0_saved_model/saved_model.pb",
+        "testdata/tftrt_2.0_saved_model/variables/variables.data-00000-of-00002",
+        "testdata/tftrt_2.0_saved_model/variables/variables.data-00001-of-00002",
+        "testdata/tftrt_2.0_saved_model/variables/variables.index",
+    ],
+    visibility = ["//tensorflow/python/compiler/tensorrt:__pkg__"],
+)
+
+base_tags = [
+    "no_cuda_on_cpu_tap",
+    "no_rocm",
+    "no_windows",
+    "nomac",
+    # TODO(b/303453873): Re-enable tests once TensorRT has been updated
+    "notap",
+]
+
+cuda_py_strict_test(
+    name = "base_test",
+    srcs = ["base_test.py"],
+    tags = base_tags + [
+        "no_oss",  # TODO(b/165611343): Need to address the failures for CUDA 11 in OSS build.
+        "notap",  # TODO(b/291653036): Some of the base tests are failing w/ TensorRT-8
+    ],
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/compiler/tensorrt:utils",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "conv2d_test",
+    srcs = ["conv2d_test.py"],
+    tags = base_tags + ["no_oss"],  # b/198501457
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:nn_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "combined_nms_test",
+    srcs = ["combined_nms_test.py"],
+    tags = base_tags + ["no_oss"],
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/compiler/tensorrt:utils",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:image_ops_impl",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "tf_function_test",
+    srcs = ["tf_function_test.py"],
+    tags = base_tags + [
+        "manual",  # TODO(b/231239602): re-enable once naming issue is resolved.
+        "no_oss",  # TODO(b/231239602): re-enable once naming issue is resolved.
+        "notap",  # TODO(b/231239602): re-enable once naming issue is resolved.
+    ],
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compiler/tensorrt:trt_convert_py",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:constants",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/util:compat",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "batch_matmul_test",
+    srcs = ["batch_matmul_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "biasadd_matmul_test",
+    srcs = ["biasadd_matmul_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "binary_tensor_weight_broadcast_test",
+    srcs = ["binary_tensor_weight_broadcast_test.py"],
+    # TODO(b/297490791): Reenable after TensorRT regression has been fixed
+    tags = base_tags + ["no_oss"],
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "bool_test",
+    srcs = ["bool_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/compiler/tensorrt:utils",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "cast_test",
+    srcs = ["cast_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "concatenation_test",
+    srcs = ["concatenation_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "const_broadcast_test",
+    srcs = ["const_broadcast_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "data_dependent_shape_test",
+    srcs = ["data_dependent_shape_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "dynamic_input_shapes_test",
+    srcs = ["dynamic_input_shapes_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "identity_output_test",
+    srcs = ["identity_output_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "int32_test",
+    srcs = ["int32_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "lru_cache_test",
+    srcs = ["lru_cache_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "memory_alignment_test",
+    srcs = ["memory_alignment_test.py"],
+    tags = base_tags + ["no_oss"],  # TODO(b/285588022)
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "multi_connection_neighbor_engine_test",
+    srcs = ["multi_connection_neighbor_engine_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "neighboring_engine_test",
+    srcs = ["neighboring_engine_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "quantization_test",
+    srcs = ["quantization_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "rank_two_test",
+    srcs = ["rank_two_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "reshape_transpose_test",
+    srcs = ["reshape_transpose_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "topk_test",
+    srcs = ["topk_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "trt_engine_op_shape_test",
+    srcs = ["trt_engine_op_shape_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:tag_constants",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "trt_mode_test",
+    srcs = ["trt_mode_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "unary_test",
+    srcs = ["unary_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "vgg_block_nchw_test",
+    srcs = ["vgg_block_nchw_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:nn_impl",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "vgg_block_test",
+    srcs = ["vgg_block_test.py"],
+    tags = base_tags,
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        ":tf_trt_integration_test_base_srcs",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:nn",
+        "//tensorflow/python/ops:nn_impl",
+        "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+test_suite(
+    name = "tf_trt_integration_test",
+    tests = [
+        "batch_matmul_test",
+        "biasadd_matmul_test",
+        "binary_tensor_weight_broadcast_test",
+        "bool_test",
+        "cast_test",
+        "combined_nms_test",
+        "concatenation_test",
+        "const_broadcast_test",
+        "data_dependent_shape_test",
+        "dynamic_input_shapes_test",
+        "identity_output_test",
+        "int32_test",
+        "lru_cache_test",
+        "memory_alignment_test",
+        "multi_connection_neighbor_engine_test",
+        "neighboring_engine_test",
+        "quantization_test",
+        "rank_two_test",
+        "reshape_transpose_test",
+        "tf_function_test",
+        "topk_test",
+        "trt_engine_op_shape_test",
+        "trt_mode_test",
+        "unary_test",
+        "vgg_block_nchw_test",
+        "vgg_block_test",
+    ],
+)
+
+test_suite(
+    name = "tf_trt_integration_test_no_oss",
+    tests = [
+        "base_test",
+        "conv2d_test",
+    ],
+)

tensorflow/python/compiler/tensorrt/test/annotate_max_batch_sizes_test.py

@@ -0,0 +1,135 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Testing the impact of graph node _tftrt_op_max_batch_size annotation on TRTEngineOp attributes."""
+
+import unittest
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class MaxBatchSizesTestBase(trt_test.TfTrtIntegrationTestBase):
+
+  @classmethod
+  def setUpClass(cls):
+    if cls is MaxBatchSizesTestBase:
+      raise unittest.SkipTest(
+          'MaxBatchSizesTestBase defines base class for other tests.')
+    super(MaxBatchSizesTestBase, cls).setUpClass()
+
+  @property
+  def tensor_shapes(self):
+    return [[1, 512, 1, 1], [64, 2, 2, 2], [32, 4, 2, 2], [16, 8, 2, 2]]
+
+  @property
+  def max_batch_sizes(self):
+    return [shape[0] for shape in self.tensor_shapes]
+
+  def GetParams(self):
+    """Gets the build parameters for the test."""
+    return self.BuildParams(
+        self.GraphFn,
+        dtype=dtypes.float32,
+        input_shapes=[self.tensor_shapes[0]],
+        output_shapes=[self.tensor_shapes[-1]])
+
+  def ShouldRunTest(self, run_params):
+    # The maximum batch size for dynamic engines will be the actual batch size
+    # detected at runtime. Therefore, we don't run the test with dynamic
+    # engines.
+    return (not run_params.dynamic_engine, 'test static engine only.')
+
+  def GetMaxBatchSize(self, run_params):
+    """Returns the max_batch_size that the converter should use for tests."""
+    if run_params.dynamic_engine:
+      return None
+    return min(self.max_batch_sizes)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Checks that the expected engine is built.
+
+    Args:
+      run_params: the run parameters.
+
+    Returns:
+      the expected engines to build.
+
+    There shall be engines generated for each maximum batch size.
+    """
+    return [
+        f'TRTEngineOp_{seq_id:03d}'
+        for seq_id in range(len(self.max_batch_sizes))
+    ]
+
+  def ExpectedMaxBatchSizes(self, run_params):
+    """Checks that the expected maximum batch sizes for the generated engines.
+
+    Args:
+      run_params: the run parameters.
+
+    Returns:
+      the expected maximum batch sizes for the generated engines.
+
+    There shall be engines generated for each maximum batch size.
+    """
+    return self.max_batch_sizes
+
+
+class AnnotateMaxBatchSizesTest(MaxBatchSizesTestBase):
+
+  def GraphFn(self, inp):
+    """Builds a tf.Graph for the test."""
+    tensor = inp * 2.0
+    tensor = array_ops.reshape(tensor, [-1] + self.tensor_shapes[1][1:])
+    with ops.get_default_graph()._attr_scope({
+        '_tftrt_op_max_batch_size':
+            attr_value_pb2.AttrValue(i=self.max_batch_sizes[1])
+    }):
+      tensor = tensor + 3.0
+    tensor = array_ops.reshape(tensor, [-1] + self.tensor_shapes[2][1:])
+    with ops.get_default_graph()._attr_scope({
+        '_tftrt_op_max_batch_size':
+            attr_value_pb2.AttrValue(i=self.max_batch_sizes[2])
+    }):
+      tensor = tensor * 4.0
+    tensor = array_ops.reshape(tensor, [-1] + self.tensor_shapes[3][1:])
+    with ops.get_default_graph()._attr_scope({
+        '_tftrt_op_max_batch_size':
+            attr_value_pb2.AttrValue(i=self.max_batch_sizes[3])
+    }):
+      tensor += tensor + 5.0
+    return array_ops.identity(tensor, name='output_0')
+
+
+class StaticBatchSizeTest(MaxBatchSizesTestBase):
+
+  def GraphFn(self, inp):
+    """Builds a tf.Graph for the test."""
+    tensor = inp * 2.0
+    tensor = array_ops.reshape(tensor, self.tensor_shapes[1])
+    tensor = tensor + 3.0
+    tensor = array_ops.reshape(tensor, self.tensor_shapes[2])
+    tensor = tensor * 4.0
+    tensor = array_ops.reshape(tensor, self.tensor_shapes[3])
+    tensor += tensor + 5.0
+    return array_ops.identity(tensor, name='output_0')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/tensorrt/test/base_test.py

@@ -0,0 +1,319 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Basic tests for TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt import utils as trt_utils
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import test
+
+
+class SimpleSingleEngineTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing single segment."""
+    dtype = inp.dtype
+    conv_filter = constant_op.constant(
+        [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]],
+        name="weights",
+        dtype=dtype)
+    conv = nn.conv2d(
+        input=inp,
+        filter=conv_filter,
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        name="conv")
+    bias = constant_op.constant([4., 1.5, 2., 3., 5., 7.],
+                                name="bias",
+                                dtype=dtype)
+    added = nn.bias_add(conv, bias, name="bias_add")
+    relu = nn.relu(added, "relu")
+    identity = array_ops.identity(relu, "identity")
+    pool = nn_ops.max_pool(
+        identity, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
+    return array_ops.squeeze(pool, name="output_0")
+
+  def GetParams(self):
+    # TODO(aaroey): test graph with different dtypes.
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
+                            [[100, 6, 6, 6]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": [
+            "weights", "conv", "bias", "bias_add", "relu", "identity",
+            "max_pool"
+        ]
+    }
+
+
+class SimpleMultiEnginesTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing multiple segment."""
+    dtype = inp.dtype
+    conv_filter = constant_op.constant(
+        [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]],
+        name="weights",
+        dtype=dtype)
+    conv = nn.conv2d(
+        input=inp,
+        filter=conv_filter,
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        name="conv")
+    c1 = constant_op.constant(
+        np.random.randn(12, 12, 6), dtype=dtype, name="c1")
+    p = math_ops.mul(conv, c1, name="mul")
+    c2 = constant_op.constant(
+        np.random.randn(12, 12, 6), dtype=dtype, name="c2")
+    q = math_ops.div(conv, c2, name="div")
+
+    edge = self.trt_incompatible_op(q, name="incompatible")
+    one = constant_op.constant(1, name="one", dtype=dtype)
+    edge = math_ops.sub(one, edge, name="one_sub")
+    edge = math_ops.div(edge, edge, name="div1")
+    r = math_ops.add(edge, edge, name="add")
+
+    p = math_ops.sub(p, edge, name="sub")
+    q = math_ops.mul(q, edge, name="mul1")
+    s = math_ops.add(p, q, name="add1")
+    s = math_ops.sub(s, r, name="sub1")
+    return array_ops.squeeze(s, name="output_0")
+
+  def GetParams(self):
+    # TODO(aaroey): test graph with different dtypes.
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
+                            [[100, 12, 12, 6]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": [
+            "add", "add1", "c1", "div1", "mul", "mul1", "sub", "sub1", "one",
+            "one_sub"
+        ],
+        "TRTEngineOp_001": ["c2", "conv", "div", "weights"]
+    }
+
+  def setUp(self):
+    super().setUp()
+    # Disable layout optimizer, since it will convert BiasAdd with NHWC
+    # format to NCHW format under four dimensional input.
+    self.DisableNonTrtOptimizers()
+
+  def ShouldRunTest(self, run_params):
+    return (
+        trt_utils.is_linked_tensorrt_version_greater_equal(8),
+        "Test is non-hermetic with TensorRT 7",
+    )
+
+
+class SimpleMultiEnginesTest2(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing two segments."""
+    n = inp
+    for i in range(2):
+      c = constant_op.constant(float(i), name="c%d" % i)
+      n = math_ops.add(n, c, name="add%d" % i)
+      n = math_ops.mul(n, n, name="mul%d" % i)
+    n = self.trt_incompatible_op(n, name="incompatible")
+    c = constant_op.constant(2.0, name="c2")
+    n = math_ops.add(n, c, name="add2")
+    n = math_ops.mul(n, n, name="mul2")
+    c = constant_op.constant(3.0, name="c3")
+    n = math_ops.add(n, c, name="add3")
+    n = math_ops.mul(n, n, name="mul3")
+    return array_ops.squeeze(n, name="output_0")
+
+  def GetParams(self):
+    shapes = [[2, 32, 32, 3]]
+    return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
+                            output_shapes=shapes)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": ["c0", "c1", "add0", "add1", "mul0", "mul1"],
+        "TRTEngineOp_001": ["c2", "c3", "add2", "add3", "mul2", "mul3"]
+    }
+
+  def ShouldRunTest(self, run_params):
+    """Whether to run the test."""
+    # Disable the test in fp16 mode since multiple matmul and add ops together
+    # can cause overflow.
+    return (
+        (run_params.precision_mode != "FP16") and
+        not (trt_test.IsQuantizationMode(run_params.precision_mode) and
+             not run_params.use_calibration)), "test FP32 and non-calibration"
+
+
+class ConstInputTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing multiple segment."""
+    n = inp
+    c = constant_op.constant(1.0, name="c")
+    # Adds data dependency from the constant op to a trt incompatible op,
+    # and adds data dependency from the trt incompatible op to the other
+    # ops, to make sure the constant op cannot be contracted with any trt
+    # segment that depends on it.
+    n = self.trt_incompatible_binary_op(n, c, name="incompatible")
+    n = math_ops.add(n, c, name="add")
+    n = math_ops.mul(n, n, name="mul")
+    n = math_ops.add(n, n, name="add1")
+    n = self.trt_incompatible_op(n, name="incompatible1")
+    n = math_ops.add(n, c, name="add2")
+    n = math_ops.mul(n, n, name="mul1")
+    n = math_ops.add(n, n, name="add3")
+    return array_ops.squeeze(n, name="output_0")
+
+  def GetParams(self):
+    shapes = [[2, 32, 32, 3]]
+    return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
+                            output_shapes=shapes)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": ["add", "add1", "mul"],
+        "TRTEngineOp_001": ["add2", "add3", "mul1"]
+    }
+
+  def ExpectedConnections(self, run_params):
+    """Returns the expected edges."""
+    return {
+        "input_0": set(),
+        "c": set(),
+        "incompatible": {"input_0", "c"},
+        "TRTEngineOp_000": {"incompatible"},
+        "incompatible1": {"TRTEngineOp_000"},
+        "TRTEngineOp_001": {"incompatible1"},
+        "output_0": {"TRTEngineOp_001"},
+    }
+
+
+class ConstDataInputSingleEngineTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing single segment."""
+    n = inp
+    c = constant_op.constant(1.0, name="c")
+    n = math_ops.add(n, c, name="add")
+    n = math_ops.mul(n, n, name="mul")
+    n = math_ops.add(n, n, name="add1")
+    return array_ops.squeeze(n, name="output_0")
+
+  def GetParams(self):
+    shapes = [[2, 32, 32, 3]]
+    return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
+                            output_shapes=shapes)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {"TRTEngineOp_000": ["c", "add", "add1", "mul"]}
+
+
+class ConstDataInputMultipleEnginesTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing multiple segment."""
+    n = inp
+    c = constant_op.constant(1.0, name="c")
+    n = math_ops.add(n, c, name="add")
+    n = math_ops.mul(n, n, name="mul")
+    n = math_ops.add(n, n, name="add1")
+    n = self.trt_incompatible_op(n, name="incompatible1")
+    n = math_ops.add(n, c, name="add2")
+    n = math_ops.mul(n, n, name="mul1")
+    n = math_ops.add(n, n, name="add3")
+    return array_ops.squeeze(n, name="output_0")
+
+  def GetParams(self):
+    shapes = [[2, 32, 32, 3]]
+    return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
+                            output_shapes=shapes)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": ["add2", "add3", "mul1"],
+        # Why segment ["add", "add1", "mul"] was assigned segment id 1
+        # instead of 0: the parent node of this segment is actually const
+        # node 'c', but it's removed later since it's const output of the
+        # segment which is not allowed.
+        "TRTEngineOp_001": ["add", "add1", "mul"]
+    }
+
+
+class ControlDependencyTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    """Create a graph containing multiple segments."""
+    c1 = constant_op.constant(1.0, name="c1")
+    c2 = constant_op.constant(2.0, name="c2")
+    d1 = self.trt_incompatible_op(inp, name="d1")
+    d2 = self.trt_incompatible_binary_op(inp, inp, name="d2")
+    with ops.control_dependencies([d1]):
+      add = math_ops.add(inp, c1, name="add")
+    mul = math_ops.mul(add, add, name="mul")
+    add1 = math_ops.add(mul, mul, name="add1")
+    edge = self.trt_incompatible_op(add1, name="incompatible")
+    with ops.control_dependencies([d1, d2, add1]):
+      add2 = math_ops.add(edge, c2, name="add2")
+    mul1 = math_ops.mul(add2, add2, name="mul1")
+    add3 = math_ops.add(mul1, mul1, name="add3")
+    inc1 = self.trt_incompatible_binary_op(d1, add3, name="incompatible1")
+    inc2 = self.trt_incompatible_binary_op(d2, inc1, name="incompatible2")
+    return array_ops.squeeze(inc2, name="output_0")
+
+  def GetParams(self):
+    shapes = [[2, 32, 32, 3]]
+    return self.BuildParams(self.GraphFn, dtypes.float32, input_shapes=shapes,
+                            output_shapes=shapes)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": ["c1", "add", "add1", "mul"],
+        "TRTEngineOp_001": ["c2", "add2", "add3", "mul1"]
+    }
+
+  def ExpectedConnections(self, run_params):
+    """Returns the expected edges."""
+    return {
+        "input_0": set(),
+        "d1": {"input_0"},
+        "d2": {"input_0"},
+        "TRTEngineOp_000": {"input_0", "^d1"},
+        "incompatible": {"TRTEngineOp_000"},
+        "TRTEngineOp_001": {"incompatible", "^d2"},
+        "incompatible1": {"TRTEngineOp_001", "d1"},
+        "incompatible2": {"incompatible1", "d2"},
+        "output_0": {"incompatible2"},
+    }
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/batch_matmul_test.py

@@ -0,0 +1,115 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import unittest
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class BatchMatMultTestBase(trt_test.TfTrtIntegrationTestBase):
+  """Base class for BatchMatMult tests."""
+
+  # Shape inference of BatchMatMultV2 doesn't work. Use static batch size.
+  def BuildParams(self, graph_fn, dtype, input_shapes, output_shapes):
+    return self.BuildParamsWithMask(
+        graph_fn=graph_fn,
+        dtype=dtype,
+        input_shapes=input_shapes,
+        output_shapes=output_shapes,
+        input_mask=[[True] * len(s) for s in input_shapes],
+        output_mask=[[True] * len(s) for s in output_shapes],
+        extra_inputs=[],
+        extra_outputs=[])
+
+  @classmethod
+  def setUpClass(cls):
+    if cls is BatchMatMultTestBase:
+      raise unittest.SkipTest(
+          "BatchMatMultTestBase defines base class for other test.")
+    super(BatchMatMultTestBase, cls).setUpClass()
+
+
+class BatchMatMulTwoTensorTest(BatchMatMultTestBase):
+  """Testing conversion of BatchMatMul where both inputs are tensors."""
+
+  def GraphFn(self, inp, inp1):
+    x1 = math_ops.matmul(inp, inp1, name="matmul")
+    # Relu to reach minimum segment size.
+    x1 = nn.relu(x1, name="relu")
+    return array_ops.identity(x1, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32,
+                            [[12, 5, 8, 12], [12, 5, 12, 7]], [[12, 5, 8, 7]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {"TRTEngineOp_000": ["matmul", "relu"]}
+
+
+class BatchMatMulWeightBroadcastTest(BatchMatMultTestBase):
+  """Testing BatchMatMulV2: one operand is weight and both have same rank."""
+
+  def ShouldAllowTF32Computation(self):
+    return False
+
+  def GraphFn(self, inp):
+    dtype = inp.dtype
+    b = constant_op.constant(
+        np.random.randn(1, 5, 7), dtype=dtype, name="kernel")
+    x1 = math_ops.matmul(inp, b, name="matmul")
+    return array_ops.identity(x1, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[12, 9, 5]],
+                            [[12, 9, 7]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {"TRTEngineOp_000": ["matmul", "kernel"]}
+
+
+class BatchMatMulWeightBroadcastDims2Test(BatchMatMultTestBase):
+  """Testing BatchMatMulV2: weight operand must be broadcasted."""
+
+  def ShouldAllowTF32Computation(self):
+    return False
+
+  def GraphFn(self, inp):
+    dtype = inp.dtype
+    b = constant_op.constant(np.random.randn(5, 7), dtype=dtype, name="kernel")
+    x1 = math_ops.matmul(inp, b, name="matmul")
+    return array_ops.identity(x1, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[12, 9, 5]],
+                            [[12, 9, 7]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {"TRTEngineOp_000": ["matmul", "kernel"]}
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/biasadd_matmul_test.py

@@ -0,0 +1,130 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class BiasaddMatMulTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing conversion of BiasAdd MatMul in TF-TRT conversion."""
+
+  def _ConstOp(self, shape):
+    return constant_op.constant(np.random.randn(*shape), dtype=dtypes.float32)
+
+  def GraphFn(self, x):
+    input_matrix_rows = 4
+    input_matrix_columns = 144
+
+    b = self._ConstOp((input_matrix_columns, 4))
+    x1 = math_ops.matmul(x, b)
+    b = self._ConstOp((1, 4))
+    x1 = x1 + b
+
+    b = self._ConstOp((input_matrix_rows, 144))
+    x2 = self.trt_incompatible_op(x)
+    x2 = math_ops.matmul(x2, b, transpose_a=True)
+    x2 = gen_array_ops.reshape(x2, [4, -1])
+    x2 = self.trt_incompatible_op(x2)
+
+    b = self._ConstOp((4, input_matrix_columns))
+    x3 = math_ops.matmul(x, b, transpose_b=True)
+
+    b = self._ConstOp((16, input_matrix_rows))
+    x4 = self.trt_incompatible_op(x)
+    x4 = math_ops.matmul(x4, b, transpose_b=True, transpose_a=True)
+    x4 = gen_array_ops.reshape(x4, [4, -1])
+    x4 = self.trt_incompatible_op(x4)
+
+    # Note that tf.nn.bias_add supports up to 5 dimensions.
+    b = self._ConstOp((input_matrix_columns, 48))
+    x5 = math_ops.matmul(x, b)
+    b = self._ConstOp((48,))
+    x5 = nn.bias_add(x5, b)
+    # TODO(b/154672994): Put the reshape back when the bug is fixed.
+    # x5 = gen_array_ops.reshape(x5, [4, -1])
+
+    x6 = gen_array_ops.reshape(x, [4, 24, 6])
+    b = self._ConstOp((6,))
+    x6 = nn.bias_add(x6, b, data_format="NHWC")
+    x6 = gen_array_ops.reshape(x6, [4, -1])
+
+    x7 = gen_array_ops.reshape(x, [4, 12, 4, 3])
+    b = self._ConstOp((3,))
+    x7 = nn.bias_add(x7, b, data_format="NHWC")
+    x7 = gen_array_ops.reshape(x7, [4, -1])
+
+    x8 = gen_array_ops.reshape(x, [4, 4, 3, 2, 6])
+    b = self._ConstOp((6,))
+    x8 = nn.bias_add(x8, b, data_format="NHWC")
+    x8 = gen_array_ops.reshape(x8, [4, -1])
+
+    x9 = gen_array_ops.reshape(x, [4, 12, 3, 2, 2])
+    b = self._ConstOp((12,))
+    x9 = nn.bias_add(x9, b, data_format="NCHW")
+    x9 = gen_array_ops.reshape(x9, [4, -1])
+
+    x10 = gen_array_ops.reshape(x, [4, 3, 4, 12])
+    b = self._ConstOp((3,))
+    x10 = nn.bias_add(x10, b, data_format="NCHW")
+    x10 = gen_array_ops.reshape(x10, [4, -1])
+
+    x11 = gen_array_ops.reshape(x, [4, 6, 24])
+    b = self._ConstOp((6,))
+    x11 = nn.bias_add(x11, b, data_format="NCHW")
+    x11 = gen_array_ops.reshape(x11, [4, -1])
+
+    out = array_ops.concat([x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11],
+                           axis=-1)
+    return array_ops.squeeze(out, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[4, 144]],
+                            [[4, 6680]])
+
+  def setUp(self):
+    super().setUp()
+    # Disable layout optimizer, since it will convert BiasAdd with NHWC
+    # format to NCHW format under four dimensional input.
+    self.DisableNonTrtOptimizers()
+
+  def GetMaxBatchSize(self, run_params):
+    """Returns the max_batch_size that the converter should use for tests."""
+    if run_params.dynamic_engine:
+      return None
+
+    return 4
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    if run_params.dynamic_shape:
+      # Increased conversion rate in dynamic shape mode due to a few additional
+      # conversions for MatMul, Reshape and Concat ops. This increases the size
+      # of the candidate segments and results in two more TrtEngineOps.
+      return ["TRTEngineOp_000", "TRTEngineOp_001", "TRTEngineOp_002"]
+    else:
+      return ["TRTEngineOp_000"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/binary_tensor_weight_broadcast_test.py

@@ -0,0 +1,80 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import os
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+
+
+class BinaryTensorWeightBroadcastTest(trt_test.TfTrtIntegrationTestBase):
+  """Tests for scale & elementwise layers in TF-TRT."""
+
+  def _ConstOp(self, shape):
+    return constant_op.constant(np.random.randn(*shape), dtype=dtypes.float32)
+
+  def GraphFn(self, x):
+    for weights_shape in [
+        (1,),  # scale
+        (24, 1, 1),  # scale
+        (24, 24, 20),  # scale
+        (20,),  # elementwise
+        (1, 24, 1, 1),  # elementwise
+        (1, 24, 24, 1),  # elementwise
+        (1, 24, 24, 20),  # elementwise
+        (24, 20),  # elementwise
+    ]:
+      a = self._ConstOp(weights_shape)
+      f = x + a
+      x = self.trt_incompatible_op(f)
+      a = self._ConstOp(weights_shape)
+      f = a + x
+      x = self.trt_incompatible_op(f)
+    return gen_array_ops.reshape(x, [5, -1], name="output_0")
+
+  def GetParams(self):
+    # TODO(aaroey): test graph with different dtypes.
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[10, 24, 24, 20]],
+                            [[5, 23040]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    # The final reshape op is converted only in dynamic shape mode. This op is
+    # placed into a new engine due to the preceding trt_incompatible_ops.
+    num_engines = 17 if run_params.dynamic_shape else 16
+    return [f"TRTEngineOp_{i:03d}" for i in range(num_engines)]
+
+  # TODO(b/176540862): remove this routine to disallow native segment execution
+  # for TensorRT 7+.
+  def setUp(self):
+    super().setUp()
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
+    gpus = config.list_physical_devices("GPU")
+
+    logging.info("Found the following GPUs:")
+    for gpu in gpus:
+      logging.info(f"\t- {gpu}")
+      config.set_memory_growth(gpu, True)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/bool_test.py

@@ -0,0 +1,51 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.compiler.tensorrt import utils as trt_utils
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class BoolTest(trt_test.TfTrtIntegrationTestBase):
+  """Test for boolean operations in TF-TRT."""
+
+  def GraphFn(self, x1, x2):
+    x = math_ops.logical_and(x1, x2)
+    x = math_ops.logical_or(x, x2)
+    q = math_ops.not_equal(x, x2)
+    q = math_ops.logical_not(q)
+    return array_ops.identity(q, name="output_0")
+
+  def GetParams(self):
+    shape = [2, 32, 32, 3]
+    return self.BuildParams(self.GraphFn, dtypes.bool, [shape, shape], [shape])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    reason = "Boolean ops are not implemented "
+    return (run_params.dynamic_shape, reason + "in ImplicitBatch mode") \
+        if trt_utils.is_linked_tensorrt_version_greater_equal(8, 2, 0)  \
+        else (False, reason + "for TRT < 8.2.0")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/cast_test.py

@@ -0,0 +1,60 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test conversion of graphs involving INT32 tensors and operations."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class CastInt32ToFp32Test(trt_test.TfTrtIntegrationTestBase):
+  """Tests cast to FP32 are split in FP16 mode."""
+
+  def _ConstOp(self, shape, dtype):
+    return constant_op.constant(np.random.randn(*shape), dtype=dtype)
+
+  def GraphFn(self, x):
+    b_f = self._ConstOp((1, 10), dtypes.float16)
+    x_f = math_ops.cast(x, dtypes.float16)
+    x_f = math_ops.mul(x_f, b_f)  # FP16 Multiply
+
+    x_f = math_ops.cast(x_f, dtypes.float32)
+    b_f = self._ConstOp((1, 10), dtypes.float32)
+    x_f = math_ops.add(x_f, b_f)  # FP32 Add
+
+    return array_ops.identity(x_f, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[1, 10]], [[1, 10]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    return {"TRTEngineOp_000": ["AddV2", "Cast", "Const", "Mul"]}
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-02
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-02
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/combined_nms_test.py

@@ -0,0 +1,222 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Script to test TF-TensorRT conversion of CombinedNMS op."""
+
+import os
+
+from tensorflow.python.compiler.tensorrt import utils as trt_utils
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import image_ops_impl
+from tensorflow.python.platform import test
+
+
+class CombinedNmsTest(trt_test.TfTrtIntegrationTestBase):
+  """Test for CombinedNMS op in TF-TRT."""
+
+  def setUp(self):
+    super().setUp()
+    self.num_boxes = 200
+
+  def GraphFn(self, boxes, scores):
+    max_total_size = 3
+    score_threshold = 0.1
+    iou_threshold = 0.5
+    # Shapes
+    max_total_size_tensor = constant_op.constant(
+        max_total_size, dtype=dtypes.int32, name='max_total_size')
+    iou_threshold_tensor = constant_op.constant(
+        iou_threshold, dtype=dtypes.float32, name='iou_threshold')
+    score_threshold_tensor = constant_op.constant(
+        score_threshold, dtype=dtypes.float32, name='score_threshold')
+    nms_output = image_ops_impl.combined_non_max_suppression(
+        boxes,
+        scores,
+        max_total_size_tensor,
+        max_total_size_tensor,
+        iou_threshold_tensor,
+        score_threshold_tensor,
+        name='combined_nms')
+    return [
+        array_ops.identity(output, name=('output_%d' % i))
+        for i, output in enumerate(nms_output)
+    ]
+
+  def GetParams(self):
+    # Parameters
+    q = 1
+    batch_size = 2
+    num_classes = 2
+    max_total_size = 3
+
+    boxes_shape = [batch_size, self.num_boxes, q, 4]
+    scores_shape = [batch_size, self.num_boxes, num_classes]
+    nmsed_boxes_shape = [batch_size, max_total_size, 4]
+    nmsed_scores_shape = [batch_size, max_total_size]
+    nmsed_classes_shape = [batch_size, max_total_size]
+    valid_detections_shape = [batch_size]
+    return self.BuildParams(self.GraphFn, dtypes.float32,
+                            [boxes_shape, scores_shape], [
+                                nmsed_boxes_shape, nmsed_scores_shape,
+                                nmsed_classes_shape, valid_detections_shape
+                            ])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    if not run_params.dynamic_shape:
+      return {
+          'TRTEngineOp_000': [
+              'combined_nms/CombinedNonMaxSuppression', 'max_total_size',
+              'iou_threshold', 'score_threshold'
+          ]
+      }
+    else:
+      # The CombinedNMS op is currently not converted in dynamic shape mode.
+      # This branch shall be removed once the converter is updated to handle
+      # input with dynamic shape.
+      return dict()
+
+  def ShouldRunTest(self, run_params):
+    should_run, reason = super().ShouldRunTest(run_params)
+    should_run = should_run and \
+        not trt_test.IsQuantizationMode(run_params.precision_mode)
+    reason += ' and precision != INT8'
+    # Only run for TRT 7.1.3 and above.
+    return should_run and trt_utils.is_linked_tensorrt_version_greater_equal(
+        7, 1, 3), reason + ' and >= TRT 7.1.3'
+
+
+class CombinedNmsExecuteNativeSegmentTest(CombinedNmsTest):
+
+  def setUp(self):
+    super().setUp()
+    os.environ['TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION'] = 'True'
+
+  def tearDown(self):
+    super().tearDown()
+    os.environ['TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION'] = 'False'
+
+  def GetMaxBatchSize(self, run_params):
+    """Returns the max_batch_size that the converter should use for tests."""
+    if run_params.dynamic_engine:
+      return None
+
+    # Build the engine with the allowed max_batch_size less than the actual
+    # max_batch_size, to fore the runtime to execute the native segment. This
+    # is to test that combined_non_max_suppression, which doesn't have a TF GPU
+    # implementation, can be executed natively even though the it is in the
+    # the graph for the TRTEngineOp with a GPU as a default device.
+    return super().GetMaxBatchSize(run_params) - 1
+
+  def ShouldRunTest(self, run_params):
+    should_run, reason = super().ShouldRunTest(run_params)
+    # max_batch_size is only useful for selecting static engines. As such,
+    # we shouldn't run the test for dynamic engines.
+    return should_run and \
+        not run_params.dynamic_engine, reason + ' and static engines'
+
+
+class CombinedNmsTestTopK(CombinedNmsTest):
+  """Test for CombinedNMS TopK op in TF-TRT."""
+
+  def GraphFn(self, pre_nms_boxes, pre_nms_scores, max_boxes_to_draw,
+              max_detetion_points):
+
+    iou_threshold = 0.1
+    score_threshold = 0.001
+
+    max_output_size_per_class_tensor = constant_op.constant(
+        max_detetion_points,
+        dtype=dtypes.int32,
+        name='max_output_size_per_class')
+
+    max_total_size_tensor = constant_op.constant(
+        max_boxes_to_draw, dtype=dtypes.int32, name='max_total_size')
+
+    iou_threshold_tensor = constant_op.constant(
+        iou_threshold, dtype=dtypes.float32, name='iou_threshold')
+
+    score_threshold_tensor = constant_op.constant(
+        score_threshold, dtype=dtypes.float32, name='score_threshold')
+
+    nms_output = image_ops_impl.combined_non_max_suppression(
+        pre_nms_boxes,
+        pre_nms_scores,
+        max_output_size_per_class=max_output_size_per_class_tensor,
+        max_total_size=max_total_size_tensor,
+        iou_threshold=iou_threshold_tensor,
+        score_threshold=score_threshold_tensor,
+        pad_per_class=False,
+        name='combined_nms')
+
+    return [
+        array_ops.identity(output, name=('output_%d' % i))
+        for i, output in enumerate(nms_output)
+    ]
+
+  def GetParams(self):
+
+    # Parameters
+    batch_size = 1
+    max_detetion_points = 2048
+    num_classes = 90
+    max_boxes_to_draw = 30
+
+    # Inputs
+    pre_nms_boxes_shape = [batch_size, max_detetion_points, 1, 4]
+    pre_nms_scores_shape = [batch_size, max_detetion_points, num_classes]
+
+    # Outputs
+    nmsed_boxes_shape = [batch_size, max_boxes_to_draw, 4]
+    nmsed_scores_shape = [batch_size, max_boxes_to_draw]
+    nmsed_classes_shape = [batch_size, max_boxes_to_draw]
+    valid_detections_shape = [batch_size]
+
+    def _get_graph_fn(x, y):
+      return self.GraphFn(
+          x,
+          y,
+          max_boxes_to_draw=max_boxes_to_draw,
+          max_detetion_points=max_detetion_points)
+
+    return self.BuildParams(_get_graph_fn, dtypes.float32,
+                            [pre_nms_boxes_shape, pre_nms_scores_shape], [
+                                nmsed_boxes_shape, nmsed_scores_shape,
+                                nmsed_classes_shape, valid_detections_shape
+                            ])
+
+
+class CombinedNmsTopKOverride(CombinedNmsTest):
+
+  def setUp(self):
+    super().setUp()
+    self.num_boxes = 5000
+    os.environ['TF_TRT_ALLOW_NMS_TOPK_OVERRIDE'] = '1'
+
+  def tearDown(self):
+    super().tearDown()
+    os.environ['TF_TRT_ALLOW_NMS_TOPK_OVERRIDE'] = '0'
+
+  def GetMaxBatchSize(self, run_params):
+    """Returns the max_batch_size that the converter should use for tests."""
+    if run_params.dynamic_engine:
+      return None
+    return super().GetMaxBatchSize(run_params)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/tensorrt/test/concatenation_test.py

@@ -0,0 +1,73 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.platform import test
+
+
+class ConcatenationTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing Concatenation in TF-TRT conversion."""
+
+  def GraphFn(self, x):
+    dtype = x.dtype
+    # scale
+    a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
+    r1 = x / a
+    a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
+    r2 = a / x
+    a = constant_op.constant(np.random.randn(1, 3, 1), dtype=dtype)
+    r3 = a + x
+    a = constant_op.constant(np.random.randn(1, 3, 1), dtype=dtype)
+    r4 = x * a
+    a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
+    r5 = x - a
+    a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
+    r6 = a - x
+    a = constant_op.constant(np.random.randn(3, 1), dtype=dtype)
+    r7 = x - a
+    a = constant_op.constant(np.random.randn(3, 1), dtype=dtype)
+    r8 = a - x
+    a = constant_op.constant(np.random.randn(3, 1, 1), dtype=dtype)
+    r9 = gen_math_ops.maximum(x, a)
+    a = constant_op.constant(np.random.randn(3, 1), dtype=dtype)
+    r10 = gen_math_ops.minimum(a, x)
+    a = constant_op.constant(np.random.randn(3), dtype=dtype)
+    r11 = x * a
+    a = constant_op.constant(np.random.randn(1), dtype=dtype)
+    r12 = a * x
+    concat1 = array_ops.concat([r1, r2, r3, r4, r5, r6], axis=-1)
+    concat2 = array_ops.concat([r7, r8, r9, r10, r11, r12], axis=3)
+    x = array_ops.concat([concat1, concat2], axis=-1)
+    return gen_array_ops.reshape(x, [2, -1], name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 3, 3, 1]],
+                            [[2, 126]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/const_broadcast_test.py

@@ -0,0 +1,61 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Script to test TF-TensorRT integration."""
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class ConstBroadcastTest(trt_test.TfTrtIntegrationTestBase):
+  """Test for Constant broadcasting in TF-TRT."""
+
+  def GraphFn(self, x):
+    """Return the expected graph to convert."""
+    dtype = x.dtype
+    filt1 = constant_op.constant(
+        0.3, shape=(3, 3, 2, 1), dtype=dtype, name='filt1')
+    y1 = nn.conv2d(x, filt1, strides=[1, 1, 1, 1], padding='SAME', name='y1')
+    z1 = nn.relu(y1, name='z1')
+    filt2 = constant_op.constant(
+        0.3, shape=(3, 3, 1, 1), dtype=dtype, name='filt2')
+    y2 = nn.conv2d(z1, filt2, strides=[1, 1, 1, 1], padding='SAME', name='y2')
+    z2 = nn.relu(y2, name='z')
+    filt3 = constant_op.constant(
+        0.3, shape=(3, 3, 1, 1), dtype=dtype, name='filt3')
+    y3 = nn.conv2d(z2, filt3, strides=[1, 1, 1, 1], padding='SAME', name='y3')
+    return nn.relu(y3, name='output_0')
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[5, 12, 12, 2]],
+                            [[5, 12, 12, 1]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ['TRTEngineOp_000']
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 1.e-04 if run_params.precision_mode == 'FP32' else 1.e-02
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 1.e-04 if run_params.precision_mode == 'FP32' else 1.e-02
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/tensorrt/test/conv2d_test.py

@@ -0,0 +1,207 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import test
+
+
+def conv2d_layer(inputs,
+                 filters,
+                 kernel_size,
+                 strides=(1, 1),
+                 padding="valid",
+                 data_format="channels_last",
+                 dilation_rate=(1, 1)):
+  dtype = inputs.dtype
+  c_axis = -1 if data_format == "channels_last" else 1
+  nchan = inputs.shape[c_axis]
+  weights_shape = (kernel_size[0], kernel_size[1], nchan, filters)
+  weights = constant_op.constant(np.random.randn(*weights_shape), dtype=dtype)
+  padding = padding.upper()
+  if data_format == "channels_last":
+    strides = [1] + list(strides) + [1]
+    dilations = [1] + list(dilation_rate) + [1]
+    data_format = "NHWC"
+  else:
+    strides = [1, 1] + list(strides)
+    dilations = [1, 1] + list(dilation_rate)
+    data_format = "NCHW"
+  return gen_nn_ops.conv2d(
+      inputs,
+      weights,
+      strides=strides,
+      padding=padding,
+      dilations=dilations,
+      data_format=data_format)
+
+
+def div_round_up(n, d):
+  return (n - 1) // d + 1
+
+
+def build_graph(inp,
+                dtype,
+                num_filters,
+                data_format,
+                kernel_sizes,
+                dilation_rates,
+                padding="same"):
+  results = []
+  for kernel_size in kernel_sizes:
+    for dilation_rate in dilation_rates:
+      result = conv2d_layer(inp, num_filters, kernel_size, (1, 1), padding,
+                            data_format, dilation_rate)
+      results.append(result)
+  output = sum(results)
+  return array_ops.identity(output, name="output_0")
+
+
+class Conv2DNCHWTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion."""
+
+  def GraphFn(self, inp):
+    np.random.seed(1234)
+    return build_graph(
+        inp=inp,
+        dtype=dtypes.float32,
+        num_filters=5,
+        data_format="channels_first",
+        kernel_sizes=[(3, 3), (3, 2)],
+        dilation_rates=[(1, 1), (2, 3)])
+
+  def GetParams(self):
+    # TODO(aaroey): test graph with different dtypes.
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 3, 7, 11]],
+                            [[13, 5, 7, 11]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    if trt_test.IsQuantizationWithCalibration(run_params):
+      return 4e-02
+    return super(Conv2DNCHWTest, self).ExpectedAbsoluteTolerance(run_params)
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    if trt_test.IsQuantizationWithCalibration(run_params):
+      return 4e-02
+    return super(Conv2DNCHWTest, self).ExpectedRelativeTolerance(run_params)
+
+
+class Conv2DNHWCTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion."""
+
+  def GraphFn(self, inp):
+    np.random.seed(1234)
+    return build_graph(
+        inp=inp,
+        dtype=dtypes.float32,
+        num_filters=5,
+        data_format="channels_last",
+        kernel_sizes=[(3, 3), (3, 2)],
+        dilation_rates=[(1, 1), (2, 3)])
+
+  def GetParams(self):
+    # TODO(aaroey): test graph with different dtypes.
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 7, 11, 3]],
+                            [[13, 7, 11, 5]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+
+class Conv2DStridedNCHWTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing conversion of strided Conv2D (data_format=NCHW)."""
+
+  def GraphFn(self, inp):
+    np.random.seed(1234)
+    num_filters = 5
+    output = inp
+    output = conv2d_layer(
+        output,
+        num_filters, (3, 2),
+        strides=(2, 2),
+        padding="same",
+        data_format="channels_first")
+    output = conv2d_layer(
+        output,
+        num_filters, (3, 3),
+        strides=(2, 2),
+        dilation_rate=(2, 3),
+        padding="same",
+        data_format="channels_first")
+    return array_ops.identity(output, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 3, 7, 11]],
+                            [[13, 5, 2, 3]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 5.e-01 if run_params.precision_mode == "INT8" else 1.e-02
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 5.e-00 if run_params.precision_mode == "INT8" else 1.e-02
+
+
+class Conv2DTranposeTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing conversion of conv2d_transpose (AKA Conv2DBackpropInput)"""
+
+  def GraphFn(self, inp):
+    np.random.seed(1234)
+    dtype = inp.dtype
+    n, c, h, w = 13, 3, 7, 11
+    num_filters = 8
+    weights_shape = [2, 2, num_filters, c]
+    weights = constant_op.constant(np.random.randn(*weights_shape), dtype=dtype)
+    output_shape = constant_op.constant([n, num_filters, h * 2, w * 2],
+                                        dtype=dtypes.int32)
+    output = nn_ops.conv2d_transpose(
+        inp,
+        weights,
+        output_shape,
+        strides=[1, 1, 2, 2],
+        padding="SAME",
+        data_format="NCHW")
+    return array_ops.identity(output, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[13, 3, 7, 11]],
+                            [[13, 8, 14, 22]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/data_dependent_shape_test.py

@@ -0,0 +1,74 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration with data dependent shapes"""
+
+import os
+
+from unittest import SkipTest  # pylint: disable=g-importing-member
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class TrtModeTestBase(trt_test.TfTrtIntegrationTestBase):
+  """Creates a network that has data dependent shapes."""
+
+  def GraphFn(self, x):
+    # With the unique() op we create a tensor with data dependent shape.
+    x = math_ops.floor(x * 10)
+    y, idx = array_ops.unique(x)
+    y = y * 2 + y
+
+    # The rest is only needed to ensure that the output has the same size as
+    # the input (expected by the test harness).
+    padding = array_ops.constant([0])
+    n = array_ops.shape(x) - array_ops.shape(y)
+    padding = array_ops.concat([padding, n], 0)
+    padding = array_ops.expand_dims(padding, 0)
+    y = array_ops.pad(y, padding)
+
+    return array_ops.identity(y, name="output_0")
+
+  def ShouldRunTest(self, run_params):
+    # We have a single dimension, and that is changed, therefore the graph can
+    # only be converted in dynamic shape mode.
+    return (run_params.dynamic_engine and run_params.is_v2 and
+            run_params.dynamic_shape and
+            run_params.use_calibration, "test v2 dynamic engine and "
+            "calibration")
+
+  def setUp(self):
+    super().setUp()
+    # The input shape depends on random input data. It can happen that we do
+    # not have engine for the actual shape. Therefore we enable native segment
+    # execution.
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
+
+  def tearDown(self):
+    super().tearDown()
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "False"
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[12]], [[12]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    return ["TRTEngineOp_000", "TRTEngineOp_001"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/dynamic_input_shapes_test.py

@@ -0,0 +1,102 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Script to test TF-TRT INT8 conversion without calibration on Mnist model."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class DynamicInputShapesTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, x):
+    conv_filter1 = constant_op.constant(
+        np.ones([3, 3, 1, 8]), name="weights1", dtype=dtypes.float32)
+    bias1 = constant_op.constant(np.random.randn(8), dtype=dtypes.float32)
+    x = nn.conv2d(
+        input=x,
+        filter=conv_filter1,
+        strides=[1, 1, 1, 1],
+        padding="SAME",
+        name="conv")
+    x = nn.bias_add(x, bias1)
+    x = nn.relu(x)
+    conv_filter2 = constant_op.constant(
+        np.ones([3, 3, 8, 1]), name="weights2", dtype=dtypes.float32)
+    bias2 = constant_op.constant(np.random.randn(1), dtype=dtypes.float32)
+    x = nn.conv2d(
+        input=x,
+        filter=conv_filter2,
+        strides=[1, 1, 1, 1],
+        padding="SAME",
+        name="conv")
+    x = nn.bias_add(x, bias2)
+    return array_ops.identity(x, name="output")
+
+  def GetParams(self):
+    # TODO(laigd): we should test the following cases:
+    # - batch size is not changed, other dims are changing
+    # - batch size is decreasing, other dims are identical
+    # - batch size is decreasing, other dims are changing
+    # - batch size is increasing, other dims are identical
+    # - batch size is increasing, other dims are changing
+    input_dims = [[[1, 5, 5, 1]], [[10, 5, 5, 1]], [[3, 5, 5, 1]],
+                  [[1, 5, 5, 1]], [[1, 3, 1, 1]], [[2, 9, 9, 1]],
+                  [[1, 224, 224, 1]], [[1, 128, 224, 1]]]
+    expected_output_dims = input_dims
+
+    return trt_test.TfTrtIntegrationTestParams(
+        graph_fn=self.GraphFn,
+        input_specs=[
+            tensor_spec.TensorSpec([None, None, None, 1], dtypes.float32,
+                                   "input")
+        ],
+        output_specs=[
+            tensor_spec.TensorSpec([None, None, None, 1], dtypes.float32,
+                                   "output")
+        ],
+        input_dims=input_dims,
+        expected_output_dims=expected_output_dims)
+
+  def setUp(self):
+    super().setUp()
+    # Disable layout optimizer, since it will convert BiasAdd with NHWC
+    # format to NCHW format under four dimensional input.
+    self.DisableNonTrtOptimizers()
+
+  def ExpectedEnginesToBuild(self, run_params):
+    return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    return (run_params.dynamic_engine and not trt_test.IsQuantizationMode(
+        run_params.precision_mode)), "test dynamic engine and non-INT8"
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-01
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-01
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/identity_output_test.py

@@ -0,0 +1,52 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""This test checks a situation where the same tensor is considered as an output
+
+multiple times because it has been duplicated by 2+ identity ops. Previously,
+the tensor would be renamed multiple times, overwriting the output binding name
+which resulted in a runtime error when the binding would not be found.
+"""
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class IdentityTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing engine with the same tensor repeated as output via identity."""
+
+  def GraphFn(self, x):
+    x1 = math_ops.exp(x)
+    x1 = x1 + x
+
+    out1 = array_ops.identity(x1, name='output_0')
+    out2 = array_ops.identity(x1, name='output_1')
+    iden1 = array_ops.identity(x1)
+    out3 = array_ops.identity(iden1, name='output_2')
+    return [out1, out2, out3]
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 32]],
+                            [[100, 32]] * 3)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ['TRTEngineOp_000']
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/tensorrt/test/int32_test.py

@@ -0,0 +1,86 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test conversion of graphs involving INT32 tensors and operations."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class ExcludeUnsupportedInt32Test(trt_test.TfTrtIntegrationTestBase):
+  """Test exclusion of ops which are not supported in INT32 mode by TF-TRT"""
+
+  def _ConstOp(self, shape, dtype):
+    return constant_op.constant(np.random.randn(*shape), dtype=dtype)
+
+  def GraphFn(self, x):
+    dtype = x.dtype
+    b = self._ConstOp((4, 10), dtype)
+    x = math_ops.matmul(x, b)
+    b = self._ConstOp((10,), dtype)
+    x = nn.bias_add(x, b)
+    return array_ops.identity(x, name='output_0')
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.int32, [[100, 4]], [[100, 10]])
+
+  def setUp(self):
+    super().setUp()
+    # Disable layout optimizer, since it will convert BiasAdd with NHWC
+    # format to NCHW format under four dimensional input.
+    self.DisableNonTrtOptimizers()
+
+  def GetMaxBatchSize(self, run_params):
+    """Returns the max_batch_size that the converter should use for tests."""
+    if run_params.dynamic_engine:
+      return None
+    return 100
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return []
+
+
+class CalibrationInt32Support(trt_test.TfTrtIntegrationTestBase):
+  """Test execution of calibration with int32 input"""
+
+  def GraphFn(self, inp):
+    # Can use any op that is converted to TRT with int32 inputs
+    inp_transposed = array_ops.transpose(inp, [0, 3, 2, 1], name='transpose_0')
+    return array_ops.identity(inp_transposed, name='output_0')
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.int32, [[3, 4, 5, 6]],
+                            [[3, 6, 5, 4]])
+
+  def ShouldRunTest(self, run_params):
+    # Although test passes with all configurations but only
+    # execute INT8 with use_calibration=True because
+    # that is the purpose of the test.
+    return trt_test.IsQuantizationWithCalibration(
+        run_params), 'test calibration and INT8'
+
+  def ExpectedEnginesToBuild(self, run_params):
+    return ['TRTEngineOp_000']
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/tensorrt/test/lru_cache_test.py

@@ -0,0 +1,66 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test LRUCache by running different input batch sizes on same network."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class LRUCacheTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, x):
+    bias = constant_op.constant(
+        np.random.randn(1, 10, 10, 1), dtype=dtypes.float32)
+    x = math_ops.add(x, bias)
+    x = nn.relu(x)
+    return array_ops.identity(x, name="output")
+
+  def GetParams(self):
+    dtype = dtypes.float32
+    input_dims = [[[1, 10, 10, 2]], [[2, 10, 10, 2]], [[4, 10, 10, 2]],
+                  [[2, 10, 10, 2]]]
+    expected_output_dims = [[[1, 10, 10, 2]], [[2, 10, 10, 2]], [[4, 10, 10,
+                                                                  2]],
+                            [[2, 10, 10, 2]]]
+    return trt_test.TfTrtIntegrationTestParams(
+        graph_fn=self.GraphFn,
+        input_specs=[
+            tensor_spec.TensorSpec([None, 10, 10, 2], dtypes.float32, "input")
+        ],
+        output_specs=[
+            tensor_spec.TensorSpec([None, 10, 10, 1], dtypes.float32, "output")
+        ],
+        input_dims=input_dims,
+        expected_output_dims=expected_output_dims)
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    return (run_params.dynamic_engine and not trt_test.IsQuantizationMode(
+        run_params.precision_mode)), "test dynamic engine and non-INT8"
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/memory_alignment_test.py

@@ -0,0 +1,70 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+@test_util.run_all_without_tensor_float_32("Uses matmul")
+class MemoryAlignmentTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing conversion of BatchMatMul in TF-TRT conversion."""
+
+  def GraphFn(self, inp):
+    dtype = inp.dtype
+    e1 = constant_op.constant(
+        np.random.randn(1, 1, 3, 5), name="kernel_1", dtype=dtype)
+    e2 = constant_op.constant(
+        np.random.randn(1, 1, 5, 10), name="kernel_2", dtype=dtype)
+    conv = nn.conv2d(
+        input=inp,
+        filter=e1,
+        strides=[1, 1, 1, 1],
+        padding="VALID",
+        name="conv")
+    out = nn.conv2d(
+        input=conv,
+        filter=e2,
+        strides=[1, 1, 1, 1],
+        padding="VALID",
+        name="conv_2")
+    return array_ops.squeeze(out, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 15, 15, 3]],
+                            [[2, 15, 15, 10]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 1.e-06 if run_params.precision_mode == "FP32" else 1.e-02
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 0.1
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/multi_connection_neighbor_engine_test.py

@@ -0,0 +1,73 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class MultiConnectionNeighborEngineTest(trt_test.TfTrtIntegrationTestBase):
+  """Test for multi connection neighboring nodes wiring tests in TF-TRT."""
+
+  def GraphFn(self, x):
+    dtype = x.dtype
+    e = constant_op.constant(
+        np.random.normal(.05, .005, [3, 2, 3, 4]), name="weights", dtype=dtype)
+    conv = nn.conv2d(
+        input=x,
+        filter=e,
+        data_format="NCHW",
+        strides=[1, 1, 1, 1],
+        padding="VALID",
+        name="conv")
+    b = constant_op.constant(
+        np.random.normal(2.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
+    t = conv + b
+
+    b = constant_op.constant(
+        np.random.normal(5.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
+    q = conv - b
+    edge = self.trt_incompatible_op(q)
+
+    b = constant_op.constant(
+        np.random.normal(5.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
+    d = b + conv
+    edge3 = self.trt_incompatible_op(d)
+
+    edge1 = self.trt_incompatible_op(conv)
+    t = t - edge1
+    q = q + edge
+    t = t + q
+    t = t + d
+    t = t - edge3
+    return array_ops.squeeze(t, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 3, 7, 5]],
+                            [[2, 4, 5, 4]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000", "TRTEngineOp_001"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/neighboring_engine_test.py

@@ -0,0 +1,62 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+
+class NeighboringEngineTest(trt_test.TfTrtIntegrationTestBase):
+  """Neighboring node wiring tests in TF-TRT conversion."""
+
+  def GraphFn(self, x):
+    dtype = x.dtype
+    e = constant_op.constant(
+        np.random.normal(.3, 0.05, [3, 2, 3, 4]), name="weights", dtype=dtype)
+    conv = nn.conv2d(
+        input=x,
+        filter=e,
+        data_format="NCHW",
+        strides=[1, 1, 1, 1],
+        padding="VALID",
+        name="conv")
+    b = constant_op.constant(
+        np.random.normal(1.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype)
+    t = math_ops.mul(conv, b, name="mul")
+    e = self.trt_incompatible_op(conv, name="incompatible")
+    t = math_ops.sub(t, e, name="sub")
+    return array_ops.squeeze(t, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[2, 3, 7, 5]],
+                            [[2, 4, 5, 4]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": ["bias", "mul", "sub"],
+        "TRTEngineOp_001": ["weights", "conv"]
+    }
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/quantization_test.py

@@ -0,0 +1,133 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def _GraphFn(x, add_quantization_nodes):
+
+  def _Quantize(x, r):
+    if add_quantization_nodes:
+      x = gen_array_ops.fake_quant_with_min_max_vars(x, -r, r)
+    return x
+
+  x = _Quantize(x, 10.0)
+  x = x + 5
+  x = _Quantize(x, 15.0)
+  x = x - 5
+  x = _Quantize(x, 10.0)
+  x = x * 0.1
+  x = _Quantize(x, 1.0)
+  w = constant_op.constant(np.ones((8, 1)), dtype=dtypes.float32)
+  x = math_ops.matmul(x, w)
+  x = _Quantize(x, 10.0)
+  return array_ops.identity(x, name="output_0")
+
+
+def _GetParams(self):
+  return self.BuildParams(self.GraphFn, dtypes.float32, [[8, 8]], [[8, 1]])
+
+
+class QuantizationMissingAllRangesTest(trt_test.TfTrtIntegrationTestBase):
+  """Create a graph containing single segment with no quantization ranges."""
+
+  def GraphFn(self, x):
+    return _GraphFn(x, add_quantization_nodes=False)
+
+  def GetParams(self):
+    return _GetParams(self)
+
+  def ShouldRunTest(self, run_params):
+    # Only test static engine mode, with or without calibration.
+    return (trt_test.IsQuantizationMode(run_params.precision_mode) and
+            not run_params.convert_online and not run_params.dynamic_engine
+           ), "test static engine, offline conversion and INT8"
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    # In static engine mode with calibration, it should build a calibration
+    # engine.
+    # In static engine mode without calibration, the engine building will
+    # succeed but fall back to non-quantized ops.
+    return ["TRTEngineOp_000"]
+
+
+class QuantizationWithRangesTest(trt_test.TfTrtIntegrationTestBase):
+  """Create a graph containing single segment with no quantization ranges."""
+
+  def GraphFn(self, x):
+    return _GraphFn(x, add_quantization_nodes=True)
+
+  def GetParams(self):
+    return _GetParams(self)
+
+  def ShouldRunTest(self, run_params):
+    # Test static/dynamic engine with/without calibration.
+    return (trt_test.IsQuantizationMode(run_params.precision_mode) and
+            not run_params.convert_online), "test offline conversion and INT8"
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
+
+
+class NonQuantizedPrecisionsWithRangesTest(trt_test.TfTrtIntegrationTestBase):
+  """Create a graph containing single segment with no quantization ranges."""
+
+  def GraphFn(self, x):
+    return _GraphFn(x, add_quantization_nodes=True)
+
+  def GetParams(self):
+    return _GetParams(self)
+
+  def ShouldRunTest(self, run_params):
+    # Only test FP32/FP16 mode.
+    return not trt_test.IsQuantizationMode(
+        run_params.precision_mode), "test non-INT8"
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    # The fake quant ops are not supported in FP32/FP16 mode, and will split the
+    # graph into two TRT segments.
+    return ["TRTEngineOp_000", "TRTEngineOp_001"]
+
+  def ExpectedAbsoluteTolerance(self, run_params):
+    """The absolute tolerance to compare floating point results."""
+    return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
+
+  def ExpectedRelativeTolerance(self, run_params):
+    """The relative tolerance to compare floating point results."""
+    return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/rank_two_test.py

@@ -0,0 +1,85 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class RankTwoTest(trt_test.TfTrtIntegrationTestBase):
+  """Test for rank 2 input in TF-TRT."""
+
+  def GraphFn(self, x1, x2):
+    # Two paths: first with rank 2 input, second with rank 4 input.
+    outputs = []
+    xs = [x1, x2]
+    for i in range(2):
+      x = xs[i]
+      c = constant_op.constant(1.0, name="c%d_1" % i)
+      q = math_ops.add(x, c, name="add%d_1" % i)
+      q = math_ops.abs(q, name="abs%d_1" % i)
+      c = constant_op.constant(2.2, name="c%d_2" % i)
+      q = math_ops.add(q, c, name="add%d_2" % i)
+      q = math_ops.abs(q, name="abs%d_2" % i)
+      c = constant_op.constant(3.0, name="c%d_3" % i)
+      q = math_ops.add(q, c, name="add%d_3" % i)
+      if i == 0:
+        axis = constant_op.constant(-1, dtype=dtypes.int32, name="axis")
+        for j in range(2):
+          q = array_ops.expand_dims(q, axis, name="expand%d_%d" % (i, j))
+        q = self.trt_incompatible_op(q)
+      q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i)
+      outputs.append(q)
+    # Combine both paths
+    q = math_ops.add(outputs[0], outputs[1], name="add")
+    return array_ops.squeeze(q, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32,
+                            [[12, 5], [12, 5, 2, 2]], [[12, 5, 2, 2]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    expected_engines = {
+        "TRTEngineOp_000": [
+            "add0_1", "add0_2", "add0_3", "c0_1", "c0_2", "c0_3", "abs0_1",
+            "abs0_2", "expand0_0", "expand0_1", "axis"
+        ],
+        "TRTEngineOp_001": [
+            "add1_1", "add1_2", "add1_3", "c1_1", "c1_2", "c1_3", "abs1_1",
+            "abs1_2", "reciprocal1"
+        ]
+    }
+    if not run_params.dynamic_shape:
+      # The two ops can't be in the same cluster as the ops in TRTEngineOp_000
+      # due to trt_incompatible_op. They can't be in the same cluster as the
+      # ops in TRTEngineOP_1 because their batch size belongs to a different
+      # equivalent class.
+      expected_engines["TRTEngineOp_002"] = ["add", "reciprocal0"]
+    else:
+      # In dynamic shape mode the batch size of the ops can differ,
+      # therefore the final ops will be merged to TRTEngineOP_1.
+      expected_engines["TRTEngineOp_001"] += ["add", "reciprocal0"]
+
+    return expected_engines
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/reshape_transpose_test.py

@@ -0,0 +1,126 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Basic tests for TF-TensorRT integration."""
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class ReshapeTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    outputs = []
+    # Here we test two types of reshapes, one changes the batch dimension and
+    # the other does not. Note that we're not able to test reshaping to
+    # scalar, since TRT requires input tensor to be of rank at least 2, so a
+    # reshape with scalar input will be filtered out of the segment before
+    # conversion.
+    #
+    # These reshapes happen at batch dimension, thus conversion should fail.
+    orig_shape = constant_op.constant([-1, 24, 24, 2], name="original_shape")
+    for shape in [[2, 50, 24, 24, 2], [-1, 50, 24, 24, 2], [2, 50, -1, 24, 2]]:
+      incompatible_reshape = array_ops.reshape(inp, shape)
+      reshape_back = array_ops.reshape(incompatible_reshape, orig_shape)
+      outputs.append(self.trt_incompatible_op(reshape_back))
+    # Add another block with many reshapes that don't change the batch
+    # dimension.
+    compatible_reshape = array_ops.reshape(
+        inp, [-1, 24 * 24, 2], name="reshape-0")
+    compatible_reshape = array_ops.reshape(
+        compatible_reshape, [100, 24, -1], name="reshape-1")
+    compatible_reshape = array_ops.reshape(
+        compatible_reshape, [100, 24 * 2, 24], name="reshape-2")
+    compatible_reshape = array_ops.reshape(
+        compatible_reshape, [-1, 24, 24 * 2], name="reshape-3")
+    compatible_reshape = array_ops.reshape(
+        compatible_reshape, [-1, 6, 4, 24, 2], name="reshape-4")
+    compatible_reshape = array_ops.reshape(
+        compatible_reshape, [-1, 6, 4, 6, 4, 2, 1], name="reshape-5")
+    compatible_reshape = array_ops.reshape(
+        compatible_reshape, [-1, 24, 24, 2], name="reshape-6")
+    outputs.append(self.trt_incompatible_op(compatible_reshape))
+    return math_ops.add_n(outputs, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
+                            [[100, 24, 24, 2]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": ["reshape-%d" % i for i in range(7)] +
+                           ["reshape-%d/shape" % i for i in range(7)]
+    }
+
+  def ShouldRunTest(self, run_params):
+    """Whether to run the test."""
+    return (not trt_test.IsQuantizationMode(run_params.precision_mode) and
+            not run_params.dynamic_engine), "test static engine and non-INT8"
+
+
+class TransposeTest(trt_test.TfTrtIntegrationTestBase):
+
+  def GraphFn(self, inp):
+    # Add a block with compatible transposes.
+    compatible_transpose = array_ops.transpose(
+        inp, [0, 3, 1, 2], name="transpose-1")
+    compatible_transpose = array_ops.transpose(
+        compatible_transpose, [0, 2, 3, 1], name="transposeback")
+    return array_ops.identity(compatible_transpose, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
+                            [[100, 24, 24, 2]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": [
+            "transpose-1", "transpose-1/perm", "transposeback",
+            "transposeback/perm"
+        ]
+    }
+
+  def ShouldRunTest(self, run_params):
+    """Whether to run the test."""
+    return (not trt_test.IsQuantizationMode(run_params.precision_mode) and
+            not run_params.dynamic_engine), "test static engine and non-INT8"
+
+
+class IncompatibleTransposeTest(TransposeTest):
+
+  def GraphFn(self, inp):
+    # Add a block with incompatible transposes.
+    incompatible_transpose = array_ops.transpose(
+        inp, [2, 1, 0, 3], name="transpose-2")
+    excluded_transpose = array_ops.transpose(
+        incompatible_transpose, [0, 2, 3, 1], name="transpose-3")
+    return array_ops.identity(excluded_transpose, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 24, 24, 2]],
+                            [[24, 100, 2, 24]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return []
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/shape_output_test.py

@@ -0,0 +1,282 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class ShapeOutputTest(trt_test.TfTrtIntegrationTestBase):
+  """Test shape value output with TF-TRT."""
+
+  def setUp(self):
+    super().setUp()
+    self.DisableNonTrtOptimizers()
+
+  def GraphFn(self, x):
+    # The first engine returns the shape of q, which equals the shape of x. The
+    # values of x are actually not needed for the TRT engine. This way x is
+    # neither a shape tensor nor an execution tensor, we still need to set its
+    # shape (binding dimensions). We confirm with this test that the binding
+    # dimensions of x are correctly set before we execute the engine.
+    q = 2 * x + 1
+    q = array_ops.shape(q)
+    q = math_ops.cast(q, dtypes.float32)
+    q = self.trt_incompatible_op(q)
+    q = q * 2 + q * q
+    return array_ops.identity(q, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[2, 2, 5, 3]], [[4]],
+        extra_inputs=[[[8, 2, 5, 3]]],
+        extra_outputs=[[[4]]],
+        input_mask=[[False, True, True, True]],
+        output_mask=[[True]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    if run_params.dynamic_shape:
+      return ["TRTEngineOp_000", "TRTEngineOp_001"]
+    else:
+      # Second segment not converted in implicit batch mode, because its
+      # tensors have only one dimensions
+      return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    # We cannot calibrate without building the engine, we turn of INT8 test.
+    return (run_params.dynamic_shape and
+            run_params.precision_mode != "INT8", "no calibration dynamic shape")
+
+
+class ShapeOutputWithSingleInputProfile(ShapeOutputTest):
+  """Same as the previous test, but with a single input profile."""
+
+  def setUp(self):
+    super().setUp()
+    self.DisableNonTrtOptimizers()
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[2, 2, 5, 3]], [[4]],
+        extra_inputs=[],
+        extra_outputs=[],
+        input_mask=[[False, True, True, True]],
+        output_mask=[[True]])
+
+
+class ShapeOutputWithSingleInputAndReshape(trt_test.TfTrtIntegrationTestBase):
+  """Similar to the previous test, but the ShapeOp output is reshaped to 2D.
+
+  This makes the output tensor not compatible with shape tensor.
+  """
+
+  def setUp(self):
+    super().setUp()
+    self.DisableNonTrtOptimizers()
+
+  def GraphFn(self, x):
+    q = 2 * x + 1
+    q = array_ops.shape(q)
+    q = gen_array_ops.reshape(q, [2, 2])
+    q = math_ops.cast(q, dtypes.float32)
+    q = self.trt_incompatible_op(q)
+    q = q * 2 + q * q
+    return array_ops.identity(q, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[2, 2, 5, 3]], [[2, 2]],
+        extra_inputs=[],
+        extra_outputs=[],
+        input_mask=[[False, True, True, True]],
+        output_mask=[[True, True]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    return ["TRTEngineOp_000", "TRTEngineOp_001"]
+
+
+class PrunedInputTest(trt_test.TfTrtIntegrationTestBase):
+  """In TRT 7, an input tensor can be pruned if it is not used by the network.
+
+  This happens if only its shape is used, but the shape is already defined by
+  the optimization profile by setting min=max. (nvbugs/3153064)
+
+  After pruning, the TRT network has no input bindings.
+  """
+
+  def setUp(self):
+    super().setUp()
+    self.DisableNonTrtOptimizers()
+
+  def GraphFn(self, x):
+    q = array_ops.shape(x)
+    q = q * 2 + q * q
+    return array_ops.identity(q, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[1, 2, 5, 3]], [[4]],
+        extra_inputs=[],
+        extra_outputs=[],
+        input_mask=[[False, True, True, True]],
+        output_mask=[[True]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    # Shape op is only converted in dynamic shape mode.
+    return (run_params.dynamic_shape and
+            run_params.is_v2, "test v2 dynamic shape")
+
+
+class PrunedInputTest2(trt_test.TfTrtIntegrationTestBase):
+  """Two inputs, one of the is pruned."""
+
+  def setUp(self):
+    super().setUp()
+    self.DisableNonTrtOptimizers()
+
+  def GraphFn(self, x, y):
+    q = array_ops.shape(x)
+    z = y * y + y
+    z = gen_array_ops.reshape(z, q)
+    out_0 = array_ops.identity(q, name="output_0")
+    out_1 = array_ops.identity(z, name="output_1")
+    return (out_0, out_1)
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[1, 2, 5, 3], [2, 15]], [[4], [1, 2, 5, 3]],
+        extra_inputs=[],
+        extra_outputs=[],
+        input_mask=[[False, True, True, True], [False, True]],
+        output_mask=[[True], [False, True, True, True]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    # Shape op is only converted in dynamic shape mode.
+    return (run_params.dynamic_shape and
+            run_params.is_v2, "test v2 dynamic shape")
+
+
+class ShapeValueMaskTest(trt_test.TfTrtIntegrationTestBase):
+  """Confirm that 0D and 1D non int tensors are not treated as shape tensors."""
+
+  def setUp(self):
+    super().setUp()
+    # This is to test whether shape value mask is correctly set in case engine
+    # construction has failed.
+    os.environ["TF_TRT_ABORT_CUDA_ENGINE_BUILD"] = "True"
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
+
+  def tearDown(self):
+    super().tearDown()
+    os.environ["TF_TRT_ABORT_CUDA_ENGINE_BUILD"] = "False"
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "False"
+
+  def GraphFn(self, x, y):
+    q = 2 * x + y
+    return array_ops.identity(q, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float16, [[3], []], [[3]],
+        extra_inputs=[],
+        extra_outputs=[],
+        input_mask=[[True], []],
+        output_mask=[[True]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    if run_params.dynamic_shape:
+      return ["TRTEngineOp_000"]
+    else:
+      return []
+
+  def ShouldRunTest(self, run_params):
+    # We cannot calibrate without building the engine, we turn of INT8 test.
+    return (run_params.dynamic_shape and
+            run_params.precision_mode != "INT8", "no calibration dynamic shape")
+
+
+class InputProfile(trt_test.TfTrtIntegrationTestBase):
+  """The shape profiles has to fit values of shape tensors, but for regular
+
+  tensors the values do not matter. Here we test shape profile management with
+  an INT32 input tensor that is not a shape tensor. The extra inputs with
+  dim=10 would trigger an error if we mistakenly treat it as shape tensors.
+  """
+
+  def setUp(self):
+    super().setUp()
+
+    self.DisableNonTrtOptimizers()
+
+  def GraphFn(self, x):
+    z = x * x + x + 1
+    z = array_ops.identity(z, name="output_0")
+    return z
+
+  def GetParams(self):
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.int32,
+        [[4]],
+        [[4]],
+        extra_inputs=[[[5]], [[10]]],
+        extra_outputs=[[[5]], [[10]]],
+        input_mask=[[False]],
+        output_mask=[[False]],
+    )
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Returns the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  def ShouldRunTest(self, run_params):
+    # Shape op is only converted in dynamic shape mode.
+    return (
+        run_params.dynamic_shape
+        and run_params.is_v2
+        and not trt_test.IsQuantizationMode(run_params.precision_mode),
+        "Test v2 dynamic_shapes without INT8",
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/test_utils.py

@@ -0,0 +1,49 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities to test TF-TensorRT integration."""
+
+import os
+
+from contextlib import contextmanager
+
+
+@contextmanager
+def experimental_feature_scope(feature_name):
+  """Creates a context manager to enable the given experimental feature.
+
+  This helper function creates a context manager setting up an experimental
+  feature temporarily.
+
+  Example:
+
+  ```python
+  with self._experimental_feature_scope("feature_1"):
+    do_smthg()
+  ```
+
+  Args:
+    feature_name: Name of the feature being tested for activation.
+  """
+
+  env_varname = "TF_TRT_EXPERIMENTAL_FEATURES"
+  env_value_bckp = os.environ.get(env_varname, default="")
+
+  exp_features = env_value_bckp.split(",")
+  os.environ[env_varname] = ",".join(list(set(exp_features + [feature_name])))
+
+  try:
+    yield
+  finally:
+    os.environ[env_varname] = env_value_bckp

tensorflow/python/compiler/tensorrt/test/testdata/gen_tf_readvariableop_model.py

@@ -0,0 +1,61 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Saves a TensorFlow model containing ReadVariableOp nodes.
+
+   The saved model is loaded and executed by tests to check that TF-TRT can
+   successfully convert and execute models with variables without freezing.
+"""
+
+import numpy as np
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.module import module
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.saved_model import save
+
+
+class MyModel(module.Module):
+  """Simple model with two variables."""
+
+  def __init__(self):
+    self.var1 = variables.Variable(
+        np.array([[[13.]]], dtype=np.float32), name="var1")
+    self.var2 = variables.Variable(
+        np.array([[[37.]]], dtype=np.float32), name="var2")
+
+  @def_function.function
+  def __call__(self, input1, input2):
+    mul1 = input1 * self.var1
+    mul2 = input2 * self.var2
+    add = mul1 + mul2
+    sub = add - 45.
+    return array_ops.identity(sub, name="output")
+
+
+def GenerateModelWithReadVariableOp(tf_saved_model_dir):
+  """Generate a model with ReadVariableOp nodes."""
+  my_model = MyModel()
+  cfunc = my_model.__call__.get_concrete_function(
+      tensor_spec.TensorSpec([None, 1, 1], dtypes.float32),
+      tensor_spec.TensorSpec([None, 1, 1], dtypes.float32))
+  # pylint: disable=not-callable
+  save(my_model, tf_saved_model_dir, signatures=cfunc)
+
+
+if __name__ == "__main__":
+  GenerateModelWithReadVariableOp(
+      tf_saved_model_dir="tf_readvariableop_saved_model")

tensorflow/python/compiler/tensorrt/test/testdata/gen_tf_variablev2_model.py

@@ -0,0 +1,90 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Saves a TensorFlow model containing VariableV2 nodes.
+
+   The saved model is loaded and executed by tests to check that TF-TRT can
+   successfully convert and execute models with VariableV2 without freezing.
+"""
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.saved_model import builder
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.saved_model import utils
+
+ops.disable_eager_execution()
+variable_scope.disable_resource_variables()
+
+
+def GenerateModelWithVariableV2(tf_saved_model_dir):
+  """Generate a model with a VariableV2 node using TFv1 API."""
+
+  def SimpleModel():
+    """Define model with a TF graph."""
+
+    def GraphFn():
+      input1 = array_ops.placeholder(
+          dtype=dtypes.float32, shape=[None, 1, 1], name="input1")
+      input2 = array_ops.placeholder(
+          dtype=dtypes.float32, shape=[None, 1, 1], name="input2")
+      var1 = variable_scope.get_variable(
+          "var1",
+          shape=[1, 1, 1],
+          initializer=init_ops.constant_initializer([[[13.]]]),
+          dtype=dtypes.float32)
+      var2 = variable_scope.get_variable(
+          "var2",
+          shape=[1, 1, 1],
+          initializer=init_ops.constant_initializer([[[37.]]]),
+          dtype=dtypes.float32)
+      mul1 = input1 * var1
+      mul2 = input2 * var2
+      add = mul1 + mul2
+      sub = add - 45.
+      out = array_ops.identity(sub, name="output")
+      return g, input1, input2, out
+
+    g = ops.Graph()
+    with g.as_default():
+      return GraphFn()
+
+  g, input1, input2, out = SimpleModel()
+  signature_def = signature_def_utils.build_signature_def(
+      inputs={
+          "input1": utils.build_tensor_info(input1),
+          "input2": utils.build_tensor_info(input2)
+      },
+      outputs={"output": utils.build_tensor_info(out)},
+      method_name=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY)
+  saved_model_builder = builder.SavedModelBuilder(tf_saved_model_dir)
+  with session.Session(graph=g) as sess:
+    variables.global_variables_initializer().run()
+    saved_model_builder.add_meta_graph_and_variables(
+        sess, [tag_constants.SERVING],
+        signature_def_map={
+            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
+        })
+  saved_model_builder.save()
+
+
+if __name__ == "__main__":
+  GenerateModelWithVariableV2(tf_saved_model_dir="tf_variablev2_saved_model")

tensorflow/python/compiler/tensorrt/test/testdata/gen_tftrt_model.py

@@ -0,0 +1,141 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Saves a SavedModel after TensorRT conversion.
+
+   The saved model is loaded and executed by tests to ensure backward
+   compatibility across TF versions.
+   The script may not work in TF1.x.
+
+   Instructions on how to use this script:
+   - Execute the script as follows:
+       python gen_tftrt_model
+   - Rename tftrt_saved_model to what makes sense for your test.
+   - Delete directory tf_saved_model unless you want to use it.
+"""
+
+from tensorflow.python import Session
+from tensorflow.python.compiler.tensorrt import trt_convert
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.saved_model import builder
+from tensorflow.python.saved_model import save
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.saved_model import utils
+from tensorflow.python.trackable import autotrackable
+
+
+def GetGraph(input1, input2, var):
+  """Define graph."""
+  add = input1 + var
+  mul = input1 * add
+  add = mul + add
+  add = add + input2
+  out = array_ops.identity(add, name="output")
+  return out
+
+
+def GenerateModelV2(tf_saved_model_dir, tftrt_saved_model_dir):
+  """Generate and convert a model using TFv2 API."""
+
+  class SimpleModel(autotrackable.AutoTrackable):
+    """Define model with a TF function."""
+
+    def __init__(self):
+      self.v = None
+
+    @def_function.function(input_signature=[
+        tensor_spec.TensorSpec(shape=[None, 1, 1], dtype=dtypes.float32),
+        tensor_spec.TensorSpec(shape=[None, 1, 1], dtype=dtypes.float32)
+    ])
+    def run(self, input1, input2):
+      if self.v is None:
+        self.v = variables.Variable([[[1.0]]], dtype=dtypes.float32)
+      return GetGraph(input1, input2, self.v)
+
+  root = SimpleModel()
+
+  # Saved TF model
+  # pylint: disable=not-callable
+  save(
+      root,
+      tf_saved_model_dir,
+      {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: root.run})
+
+  # Convert TF model to TensorRT
+  converter = trt_convert.TrtGraphConverterV2(
+      input_saved_model_dir=tf_saved_model_dir)
+  converter.convert()
+  try:
+    line_length = max(160, os.get_terminal_size().columns)
+  except OSError:
+    line_length = 160
+  converter.summary(line_length=line_length, detailed=True)
+  converter.save(tftrt_saved_model_dir)
+
+
+def GenerateModelV1(tf_saved_model_dir, tftrt_saved_model_dir):
+  """Generate and convert a model using TFv1 API."""
+
+  def SimpleModel():
+    """Define model with a TF graph."""
+
+    def GraphFn():
+      input1 = array_ops.placeholder(
+          dtype=dtypes.float32, shape=[None, 1, 1], name="input1")
+      input2 = array_ops.placeholder(
+          dtype=dtypes.float32, shape=[None, 1, 1], name="input2")
+      var = variables.Variable([[[1.0]]], dtype=dtypes.float32, name="v1")
+      out = GetGraph(input1, input2, var)
+      return g, var, input1, input2, out
+
+    g = ops.Graph()
+    with g.as_default():
+      return GraphFn()
+
+  g, var, input1, input2, out = SimpleModel()
+  signature_def = signature_def_utils.build_signature_def(
+      inputs={
+          "input1": utils.build_tensor_info(input1),
+          "input2": utils.build_tensor_info(input2)
+      },
+      outputs={"output": utils.build_tensor_info(out)},
+      method_name=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY)
+  saved_model_builder = builder.SavedModelBuilder(tf_saved_model_dir)
+  with Session(graph=g) as sess:
+    sess.run(var.initializer)
+    saved_model_builder.add_meta_graph_and_variables(
+        sess, [tag_constants.SERVING],
+        signature_def_map={
+            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
+        })
+  saved_model_builder.save()
+
+  # Convert TF model to TensorRT
+  converter = trt_convert.TrtGraphConverter(
+      input_saved_model_dir=tf_saved_model_dir, is_dynamic_op=True)
+  converter.convert()
+  converter.save(tftrt_saved_model_dir)
+
+
+if __name__ == "__main__":
+  GenerateModelV2(
+      tf_saved_model_dir="tf_saved_model",
+      tftrt_saved_model_dir="tftrt_saved_model")

tensorflow/python/compiler/tensorrt/test/testdata/mnist/checkpoint

@@ -0,0 +1,3 @@
+model_checkpoint_path: "model.ckpt-46900"
+all_model_checkpoint_paths: "model.ckpt-0"
+all_model_checkpoint_paths: "model.ckpt-46900"

tensorflow/python/compiler/tensorrt/test/tf_function_test.py

@@ -0,0 +1,348 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""TF function conversion."""
+
+import itertools
+import os
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.python.compiler.tensorrt import trt_convert
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.compiler.tensorrt.test.tf_trt_integration_test_base import GraphState
+from tensorflow.python.compiler.tensorrt.test.tf_trt_integration_test_base import IsQuantizationWithCalibration
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import loader_impl
+from tensorflow.python.util import compat
+
+
+class TfFunctionTest(trt_test.TfTrtIntegrationTestBase):
+
+  def __init__(self, methodName):  # pylint: disable=invalid-name
+    super(TfFunctionTest, self).__init__(methodName)
+    self._profile_strategy = "Range"
+    self._trt_engine_op_count_offset = 0
+    self._test_conversion_params = {
+        "_tftrt_convert_function": True,
+        "_tftrt_trt_logger_name": "DefaultLogger",
+        "_tftrt_max_batch_size": 10,
+        "_tftrt_max_workspace_size_bytes":
+            (trt_convert.DEFAULT_TRT_MAX_WORKSPACE_SIZE_BYTES),
+        "_tftrt_precision_mode": "FP16",
+        "_tftrt_minimum_segment_size": 2,
+        "_tftrt_is_dyn_op": True,
+        "_tftrt_max_cached_engines": 1,
+        "_tftrt_use_calibration": False,
+        "_tftrt_use_implicit_batch": True,
+        "_tftrt_profile_strategy": self._profile_strategy,
+        "_tftrt_allow_build_at_runtime": False
+    }
+    self._is_v2 = False
+
+  def ShouldRunTest(self, run_params):
+    should_run, reason_for_skipping = (
+        trt_test.TfTrtIntegrationTestBase.ShouldRunTest(self, run_params))
+    if not should_run:
+      return should_run, reason_for_skipping
+    else:
+      # TODO(kyungtaek): Calibration currently does not run for nodes
+      # nested within functions. If this gets fixed, this method should not
+      # override the parent method.
+      return (not IsQuantizationWithCalibration(run_params),
+              "calibration is not supported for tf.functions")
+
+  def _copy_test_attr_to_func_def(self, func_def, param_name, attr_value_type):
+    test_value = self._test_conversion_params[param_name]
+    if attr_value_type == "s":
+      byte_value = compat.as_bytes(test_value)
+      func_def.attr[param_name].CopyFrom(attr_value_pb2.AttrValue(s=byte_value))
+    elif attr_value_type == "b":
+      func_def.attr[param_name].CopyFrom(attr_value_pb2.AttrValue(b=test_value))
+    elif attr_value_type == "i":
+      func_def.attr[param_name].CopyFrom(attr_value_pb2.AttrValue(i=test_value))
+    else:
+      logging.info("Attr_value type %s is not supported", attr_value_type)
+
+  def _ChainAllNodes(self, graph_def):
+    return itertools.chain(
+        graph_def.node,
+        itertools.chain(
+            *[func.node_def for func in graph_def.library.function]))
+
+  def _VerifyTestAttrs(self, function_protos):
+    if self._test_conversion_params["_tftrt_convert_function"]:
+      for func_def in function_protos:
+        if not func_def.signature.name.startswith("TRTEngine"):
+          for key, value in self._test_conversion_params.items():
+            self.assertIn(key, func_def.attr,
+                          "key %s not found in func_def.attr" % key)
+            if isinstance(value, str):
+              self.assertEqual(func_def.attr[key].s, compat.as_bytes(value))
+            elif isinstance(value, bool):
+              self.assertEqual(func_def.attr[key].b, value)
+            elif isinstance(value, int):
+              self.assertEqual(func_def.attr[key].i, value)
+
+  @def_function.function(input_signature=[
+      tensor_spec.TensorSpec(shape=[None, 32, 32, 2], dtype=dtypes.float32)
+  ])
+  def _conv_and_pool_0(self, inp):
+    dtype = inp.dtype
+    conv_filter = constant_op.constant([[[[1., 0.5, 4.], [1., 0.5, 1.]]]],
+                                       name="weights",
+                                       dtype=dtype)
+    conv = nn.conv2d(
+        input=inp,
+        filter=conv_filter,
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        name="conv")
+    bias = constant_op.constant([4., 1.5, 2.], name="bias", dtype=dtype)
+    added = nn.bias_add(conv, bias, name="bias_add")
+    relu = nn.relu(added, "relu")
+    identity = array_ops.identity(relu, "identity")
+    pool = nn_ops.max_pool(
+        identity, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
+    return array_ops.squeeze(pool)
+
+  def GraphFn(self, x):
+    x = self._conv_and_pool_0(x)
+    return array_ops.identity(x, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[10, 32, 32, 2]],
+                            [[10, 8, 8, 3]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {
+        "TRTEngineOp_000": [
+            "weights", "conv", "bias", "bias_add", "relu", "identity",
+            "max_pool"
+        ]
+    }
+
+  def _copy_test_attributes_to_func_def(self, func_def):
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_convert_function",
+        attr_value_type="b")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_trt_logger_name",
+        attr_value_type="s")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_max_batch_size",
+        attr_value_type="i")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_max_workspace_size_bytes",
+        attr_value_type="i")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_precision_mode",
+        attr_value_type="s")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_minimum_segment_size",
+        attr_value_type="i")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def, param_name="_tftrt_is_dyn_op", attr_value_type="b")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_max_cached_engines",
+        attr_value_type="i")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_use_calibration",
+        attr_value_type="b")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_use_implicit_batch",
+        attr_value_type="b")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_profile_strategy",
+        attr_value_type="s")
+    self._copy_test_attr_to_func_def(
+        func_def=func_def,
+        param_name="_tftrt_allow_build_at_runtime",
+        attr_value_type="b")
+
+  def _MakeSavedModelV1(self, run_params):
+    """Write the saved model as an input for testing.
+
+    In addition to creating a SavedModel like its parent method, this method
+    replaces this SavedModel by adding TF-TRT conversion parameters as function
+    attributes to each function in the SavedModel.
+
+    Args:
+      run_params: The current test run parameters.
+
+    Returns:
+      The directory of the saved model.
+    """
+    saved_model_dir = trt_test.TfTrtIntegrationTestBase._MakeSavedModelV1(
+        self, run_params)
+    saved_model_proto = loader_impl.parse_saved_model(saved_model_dir)
+    new_saved_model = saved_model_pb2.SavedModel()
+    new_saved_model.CopyFrom(saved_model_proto)
+    new_meta_graph_def = new_saved_model.meta_graphs[0]
+    for func_def in new_meta_graph_def.graph_def.library.function:
+      # Disable function inlining.
+      func_def.attr["_noinline"].CopyFrom(attr_value_pb2.AttrValue(b=True))
+      self._copy_test_attributes_to_func_def(func_def)
+    old_saved_model_file = os.path.join(saved_model_dir,
+                                        constants.SAVED_MODEL_FILENAME_PB)
+    if os.path.exists(old_saved_model_file):
+      os.remove(old_saved_model_file)
+    path = os.path.join(
+        compat.as_bytes(saved_model_dir),
+        compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+    file_io.write_string_to_file(
+        path, new_saved_model.SerializeToString(deterministic=True))
+    return saved_model_dir
+
+  def _MakeSavedModelV2(self, run_params):
+    """Write the saved model as an input for testing.
+
+    In addition to creating a SavedModel like its parent method, this method
+    replaces this SavedModel by adding TF-TRT conversion parameters as function
+    attributes to each function in the SavedModel.
+
+    Args:
+      run_params: The current test run parameters.
+
+    Returns:
+      The directory of the saved model.
+    """
+    saved_model_dir = trt_test.TfTrtIntegrationTestBase._MakeSavedModelV2(
+        self, run_params)
+    saved_model_proto = loader_impl.parse_saved_model(saved_model_dir)
+    new_saved_model = saved_model_pb2.SavedModel()
+    new_saved_model.CopyFrom(saved_model_proto)
+    new_meta_graph_def = new_saved_model.meta_graphs[0]
+    prefix_len = len("__inference_")
+    for func_def in new_meta_graph_def.graph_def.library.function:
+      logging.info("_MakeSavedModelV2, func_def name: %s",
+                   func_def.signature.name)
+      func_name_without_prefix = func_def.signature.name[prefix_len:]
+      if func_name_without_prefix.startswith(
+          ("_conv_and_pool_0")):
+        func_def.attr["_noinline"].CopyFrom(attr_value_pb2.AttrValue(b=True))
+        self._copy_test_attributes_to_func_def(func_def)
+    old_saved_model_file = os.path.join(saved_model_dir,
+                                        constants.SAVED_MODEL_FILENAME_PB)
+    if os.path.exists(old_saved_model_file):
+      os.remove(old_saved_model_file)
+    path = os.path.join(
+        compat.as_bytes(saved_model_dir),
+        compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+    file_io.write_string_to_file(
+        path, new_saved_model.SerializeToString(deterministic=True))
+    return saved_model_dir
+
+  def _VerifyGraphDefV1(self, run_params, original_gdef, gdef_to_verify,
+                        graph_state):
+    expected_engines = self.ExpectedEnginesToBuild(run_params)
+    num_engines = 0
+    functions = [f.signature.name for f in gdef_to_verify.library.function]
+    all_nodes = list(self._ChainAllNodes(gdef_to_verify))
+    all_nodes.sort(key=lambda x: x.name)
+
+    for node in all_nodes:
+      if node.op == "TRTEngineOp":
+        logging.info("Found TRTEngineOp: " + node.name)
+        num_engines += 1
+        segment_funcdef_name = node.attr["segment_func"].func.name
+        function_name = node.name + "_native_segment"
+        is_dynamic_engine = not node.attr["static_engine"].b
+        self.assertNotEmpty(segment_funcdef_name, node.name)
+        self.assertIn(function_name, functions)
+        if (not IsQuantizationWithCalibration(run_params) and
+            not is_dynamic_engine):
+          self.assertTrue(len(node.attr["serialized_segment"].s), node.name)
+        self.assertIn(
+            self._RemoveGraphSequenceNumber(node.name), expected_engines)
+        self.assertEqual(
+            self._ToBytes(run_params.precision_mode),
+            node.attr["precision_mode"].s, node.name)
+
+        self.assertEqual(run_params.dynamic_engine, is_dynamic_engine,
+                         node.name)
+        self.assertEqual(node.attr["use_calibration"].b,
+                         run_params.use_calibration, node.name)
+
+        has_calibration_data = len(node.attr["calibration_data"].s)
+        if (IsQuantizationWithCalibration(run_params) and
+            graph_state == GraphState.INFERENCE):
+          self.assertTrue(has_calibration_data, node.name)
+        else:
+          self.assertFalse(has_calibration_data, node.name)
+    if graph_state == GraphState.ORIGINAL:
+      self.assertEqual(0, num_engines)
+      self._VerifyTestAttrs(function_protos=gdef_to_verify.library.function)
+    else:
+      self.assertEqual(num_engines, len(expected_engines))
+      expected_connections = self.ExpectedConnections(run_params)
+      if expected_connections:
+        self._VerifyConnections(expected_engines, expected_connections,
+                                original_gdef, gdef_to_verify)
+      self._VerifyMaxBatchSizeAnnotations(
+          expected_engines=expected_engines,
+          original_gdef=original_gdef,
+          converted_gdef=gdef_to_verify,
+          expected_max_batch_sizes=self.ExpectedMaxBatchSizes(run_params),
+          default_max_batch_size=self.GetMaxBatchSize(run_params))
+      self._VerifyTestAttrs(function_protos=gdef_to_verify.library.function)
+
+  def _ShouldConverterBuild(self, run_params):
+    return (run_params.is_v2 and not run_params.convert_online and
+            run_params.dynamic_engine)
+
+  def RunTest(self, run_params):
+    self._test_conversion_params["_tftrt_precision_mode"] = (
+        run_params.precision_mode)
+    self._test_conversion_params["_tftrt_use_calibration"] = (
+        run_params.use_calibration)
+    self._test_conversion_params["_tftrt_is_dyn_op"] = (
+        run_params.dynamic_engine)
+    # When running with V1, using dynamic_engine and
+    # allow_build_at_runtime==False at the same time do not work.
+    if run_params.is_v2:
+      self._test_conversion_params["_tftrt_allow_build_at_runtime"] = True
+      self._is_v2 = True
+    else:
+      self._test_conversion_params["_tftrt_allow_build_at_runtime"] = (
+          run_params.convert_online or run_params.dynamic_engine)
+    self._test_conversion_params["_tftrt_use_implicit_batch"] = \
+        not run_params.dynamic_shape
+    self.DisableNonTrtOptimizers()
+    trt_test.TfTrtIntegrationTestBase.RunTest(self, run_params)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/topk_test.py

@@ -0,0 +1,72 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class TopKTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing Top-K in TF-TRT conversion."""
+
+  def GraphFn(self, x):
+    k = 5
+    k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
+    values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
+    values = array_ops.identity(values, name="output_0")
+    indices = array_ops.identity(indices, name="output_1")
+    return values, indices
+
+  def GetParams(self):
+    k = 5
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 100]],
+                            [[100, k], [100, k]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {"TRTEngineOp_000": ["Const", "TopK"]}
+
+
+class TopKOutputTypeTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing that output type of engine using Top-K is set correctly."""
+
+  def GraphFn(self, x):
+    k = 5
+    k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
+    values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
+    # Reshape will act as a layer between the TopK output and the engine
+    # output, requiring the output tensor of reshape to be set explicitly to
+    # int32.
+    indices = array_ops.reshape(indices, [100, 1, 5], name="Reshape")
+    values = array_ops.identity(values, name="output_0")
+    indices = array_ops.identity(indices, name="output_1")
+    return values, indices
+
+  def GetParams(self):
+    k = 5
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[100, 100]],
+                            [[100, k], [100, 1, k]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return {"TRTEngineOp_000": ["Const", "TopK", "Reshape", "Reshape/shape"]}
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/trt_engine_op_shape_test.py

@@ -0,0 +1,93 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""This script test input and output shapes and dtype of the TRTEngineOp."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import tag_constants
+
+
+class TRTEngineOpInputOutputShapeTest(trt_test.TfTrtIntegrationTestBase):
+  """Testing the output shape of a TRTEngine."""
+
+  def GraphFn(self, inp):
+    b = array_ops.squeeze(inp, axis=[2])
+    c = nn.relu(b)
+    d1 = c + c
+    d2 = math_ops.reduce_sum(d1)
+
+    d1 = array_ops.identity(d1, name="output_0")
+    d2 = array_ops.identity(d2, name="output_1")
+    return d1, d2
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[1, 2, 1, 4]],
+                            [[1, 2, 4], []])
+
+  def _GetInferGraph(self, *args, **kwargs):
+    trt_saved_model_dir = super(TRTEngineOpInputOutputShapeTest,
+                                self)._GetInferGraph(*args, **kwargs)
+
+    def get_func_from_saved_model(saved_model_dir):
+      try:  # Necessary for `bazel run ...`
+        saved_model_load_fn = load.load
+      except AttributeError:  # All the other cases
+        saved_model_load_fn = load
+      saved_model_loaded = saved_model_load_fn(
+          saved_model_dir, tags=[tag_constants.SERVING])
+      graph_func = saved_model_loaded.signatures[
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+      return graph_func, saved_model_loaded
+
+    func, _ = get_func_from_saved_model(trt_saved_model_dir)
+
+    input_shape = func.inputs[0].shape
+    if isinstance(input_shape, tensor_shape.TensorShape):
+      input_shape = input_shape.as_list()
+
+    output_shapes = [
+        out_shape.shape.as_list() if isinstance(
+            out_shape.shape, tensor_shape.TensorShape) else out_shape.shape
+        for out_shape in func.outputs
+    ]
+
+    self.assertEqual(func.inputs[0].dtype, dtypes.float32)
+    self.assertEqual(func.outputs[0].dtype, dtypes.float32)
+    self.assertEqual(func.outputs[1].dtype, dtypes.float32)
+
+    self.assertEqual(input_shape, [None, 2, 1, 4])
+    self.assertEqual(output_shapes[0], [None, 2, 4])
+    self.assertEqual(output_shapes[1], [])
+
+    return trt_saved_model_dir
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/trt_mode_test.py

@@ -0,0 +1,131 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+from unittest import SkipTest  # pylint: disable=g-importing-member
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class TrtModeTestBase(trt_test.TfTrtIntegrationTestBase):
+  """Test squeeze on batch dim and some unary operations in TF-TRT."""
+
+  def GraphFn(self, x1):
+    q = math_ops.abs(x1)
+    q = q + 1.0
+    q = q * 3.0
+    q = array_ops.squeeze(q, 0)
+    q = math_ops.abs(q)
+    q = q + 5.0
+    return array_ops.identity(q, name="output_0")
+
+  def ShouldRunTest(self, run_params):
+    # Squeeze op produces dynamic shaped values. Therefore, we don't run the
+    # test with static engine to avoid native segment execution.
+    return (run_params.dynamic_engine and run_params.is_v2 and
+            not run_params.use_calibration, "test v2 dynamic engine and "
+            "non-calibration")
+
+  def GetParams(self):
+    """The input has 1 as a first dimension, which is removed by the squeeze.
+
+    op in the graph.
+
+    In explicit batch mode, TensorRT can convert the whole graph. In this mode
+    it is possible to manipulate the batch dimension using the squeeze op.
+
+    In implicit batch mode TensorRT cannot convert the whole graph. We are not
+    allowed to manipulate (squeeze) the first dimension in implicit batch mode.
+    Therefore the graph will be converted using multiple segments.
+    """
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[1, 12, 5]],
+                            [[12, 5]])
+
+  def GetMaxBatchSize(self, run_params):
+    if run_params.dynamic_engine:
+      return None
+
+    # The first dimension of the input is squeezed and the batch size for the
+    # rest OPs is 12.
+    return 12
+
+  @classmethod
+  def setUpClass(cls):
+    if cls is TrtModeTestBase:
+      raise SkipTest("TrtModeTestBase defines base class for other test.")
+    super(TrtModeTestBase, cls).setUpClass()
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Check that the expected engine is built.
+
+    Args:
+      run_params: the run parameters.
+
+    Returns:
+      the expected engines to build.
+
+    The squeeze op is not converted by TensorRT in implicit batch mode.
+    Because of this we have two TRTEngineOp in the graphs: one for the
+    subgraph before 'squeeze(q,0)', and another one for the rest of the ops
+    after the 'squeeze(q,0)'.
+
+    In explicit batch mode the whole graph is converted using a single engine.
+    """
+    if run_params.dynamic_shape:
+      return ["TRTEngineOp_000"]
+    else:
+      return ["TRTEngineOp_000", "TRTEngineOp_001"]
+
+
+class StaticInputTest(TrtModeTestBase):
+
+  def GetParams(self):
+    """We specify input/output masks with static (known) shapes."""
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[1, 12, 5]], [[12, 5]],
+        input_mask=[[True, True, True]],
+        output_mask=[[True, True]],
+        extra_inputs=[],
+        extra_outputs=[])
+
+
+class DynamicInputTest(TrtModeTestBase):
+  """Test with dynamic input shapes.
+
+  The difference to the previous test is that we use input and output masks to
+  change the input and output shapes to unknown shapes.
+  """
+
+  def GetParams(self):
+    """We specify input/output mask with dynamic (unknown) shapes.
+
+    In dynamic shape mode, single engine with three optimization profiles can
+    handle the three different input shapes.
+    """
+    return self.BuildParamsWithMask(
+        self.GraphFn,
+        dtypes.float32, [[1, 12, 5]], [[12, 5]],
+        extra_inputs=[[[1, 2, 3]], [[1, 4, 6]]],
+        extra_outputs=[[[2, 3]], [[4, 6]]],
+        input_mask=[[False, False, False]],
+        output_mask=[[False, False]])
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/unary_test.py

@@ -0,0 +1,106 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class UnaryTest(trt_test.TfTrtIntegrationTestBase):
+  """Test for unary operations in TF-TRT."""
+
+  def GraphFn(self, x1, x2):
+    x = x1
+    q = math_ops.abs(x)
+    q = q + 1.0
+    q = gen_math_ops.exp(q)
+    q = gen_math_ops.log(q)
+    q = array_ops.squeeze(q, axis=-2)
+    q = math_ops.abs(q)
+    q = q + 2.2
+    q = gen_math_ops.sqrt(q)
+    q = gen_math_ops.rsqrt(q)
+    q = math_ops.negative(q)
+    q = array_ops.squeeze(q, axis=3)
+    q = math_ops.abs(q)
+    q = q + 3.0
+    a = gen_math_ops.reciprocal(q)
+
+    # this chain of operations has a batch size of 5, which is different from
+    # the batch size for the other operations.
+    x = constant_op.constant(np.random.randn(5, 8, 12), dtype=x.dtype)
+    q = math_ops.abs(x)
+    q = q + 2.0
+    q = gen_math_ops.exp(q)
+    q = gen_math_ops.log(q)
+    q = math_ops.abs(q)
+    q = q + 2.1
+    q = gen_math_ops.sqrt(q)
+    q = gen_math_ops.rsqrt(q)
+    q = math_ops.negative(q)
+    q = math_ops.abs(q)
+    q = q + 4.0
+    b = gen_math_ops.reciprocal(q)
+
+    # TODO(jie): this one will break, broadcasting on batch.
+    x = x2
+    q = math_ops.abs(x)
+    q = q + 5.0
+    q = gen_math_ops.exp(q)
+    q = array_ops.squeeze(q, axis=[-1, -2, 3])
+    q = gen_math_ops.log(q)
+    q = math_ops.abs(q)
+    q = q + 5.1
+    q = gen_array_ops.reshape(q, [12, 5, 1, 1, 8, 1, 12])
+    q = array_ops.squeeze(q, axis=[5, 2, 3])
+    q = gen_math_ops.sqrt(q)
+    q = math_ops.abs(q)
+    q = q + 5.2
+    q = gen_math_ops.rsqrt(q)
+    q = math_ops.negative(q)
+    q = math_ops.abs(q)
+    q = q + 5.3
+    c = gen_math_ops.reciprocal(q)
+
+    q = a * b
+    q = q / c
+    return array_ops.squeeze(q, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32,
+                            [[12, 5, 8, 1, 1, 12], [12, 5, 8, 1, 12, 1, 1]],
+                            [[12, 5, 8, 12]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    if run_params.dynamic_shape:
+      # All the ops are converted into a single TRTEngineOp.
+      return ["TRTEngineOp_000"]
+    else:
+      # Final squeeze and div ops are not converted. The x1 and x2 branches
+      # are segmented into separate TRTEngineOp.
+      return ["TRTEngineOp_000", "TRTEngineOp_001"]
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/vgg_block_nchw_test.py

@@ -0,0 +1,77 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import os
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import nn_impl
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import test
+
+
+class VGGBlockNCHWTest(trt_test.TfTrtIntegrationTestBase):
+  """Single vgg layer in NCHW unit tests in TF-TRT."""
+
+  def GraphFn(self, x):
+    dtype = x.dtype
+    x, _, _ = nn_impl.fused_batch_norm(
+        x, [1.0, 1.0], [0.0, 0.0],
+        mean=[0.5, 0.5],
+        variance=[1.0, 1.0],
+        data_format="NCHW",
+        is_training=False)
+    e = constant_op.constant(
+        np.random.randn(1, 1, 2, 6), name="weights", dtype=dtype)
+    conv = nn.conv2d(
+        input=x,
+        filter=e,
+        data_format="NCHW",
+        strides=[1, 1, 2, 2],
+        padding="SAME",
+        name="conv")
+    b = constant_op.constant(np.random.randn(6), name="bias", dtype=dtype)
+    t = nn.bias_add(conv, b, data_format="NCHW", name="biasAdd")
+    relu = nn.relu(t, "relu")
+    idty = array_ops.identity(relu, "ID")
+    v = nn_ops.max_pool(
+        idty, [1, 1, 2, 2], [1, 1, 2, 2],
+        "VALID",
+        data_format="NCHW",
+        name="max_pool")
+    return array_ops.squeeze(v, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[5, 2, 8, 8]],
+                            [[5, 6, 2, 2]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  # TODO(b/159459919): remove this routine to disallow native segment execution.
+  def setUp(self):
+    super().setUp()
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/test/vgg_block_test.py

@@ -0,0 +1,68 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Model script to test TF-TensorRT integration."""
+
+import os
+
+import numpy as np
+
+from tensorflow.python.compiler.tensorrt.test import tf_trt_integration_test_base as trt_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.ops import nn_impl
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import test
+
+
+class VGGBlockTest(trt_test.TfTrtIntegrationTestBase):
+  """Single vgg layer test in TF-TRT conversion."""
+
+  def GraphFn(self, x):
+    dtype = x.dtype
+    x, _, _ = nn_impl.fused_batch_norm(
+        x, [1.0, 1.0], [0.0, 0.0],
+        mean=[0.5, 0.5],
+        variance=[1.0, 1.0],
+        is_training=False)
+    e = constant_op.constant(
+        np.random.randn(1, 1, 2, 6), name="weights", dtype=dtype)
+    conv = nn.conv2d(
+        input=x, filter=e, strides=[1, 2, 2, 1], padding="SAME", name="conv")
+    b = constant_op.constant(np.random.randn(6), name="bias", dtype=dtype)
+    t = nn.bias_add(conv, b, name="biasAdd")
+    relu = nn.relu(t, "relu")
+    idty = array_ops.identity(relu, "ID")
+    v = nn_ops.max_pool(
+        idty, [1, 2, 2, 1], [1, 2, 2, 1], "VALID", name="max_pool")
+    return array_ops.squeeze(v, name="output_0")
+
+  def GetParams(self):
+    return self.BuildParams(self.GraphFn, dtypes.float32, [[5, 8, 8, 2]],
+                            [[5, 2, 2, 6]])
+
+  def ExpectedEnginesToBuild(self, run_params):
+    """Return the expected engines to build."""
+    return ["TRTEngineOp_000"]
+
+  # TODO(b/159459919): remove this routine to disallow native segment execution.
+  def setUp(self):
+    super().setUp()
+    os.environ["TF_TRT_ALLOW_ENGINE_NATIVE_SEGMENT_EXECUTION"] = "True"
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/tensorrt/utils.py

@@ -0,0 +1,259 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Exposes the Python wrapper conversion to trt_graph."""
+
+import collections
+import os
+import re
+
+from packaging import version
+
+from tensorflow.compiler.tf2tensorrt import _pywrap_py_utils
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.framework import dtypes
+
+
+def disable_non_trt_optimizers_in_rewriter_config(rewriter_config):
+  """Modifies rewriter_config to disable all non-TRT optimizations."""
+  off = rewriter_config_pb2.RewriterConfig.OFF
+
+  rewriter_config.arithmetic_optimization = off
+  rewriter_config.auto_mixed_precision = off
+  rewriter_config.auto_parallel.enable = False
+  rewriter_config.constant_folding = off
+  rewriter_config.debug_stripper = off
+  rewriter_config.dependency_optimization = off
+  # This one needs to be ON to allow TF-TRT
+  rewriter_config.disable_meta_optimizer = False
+  rewriter_config.disable_model_pruning = True
+  rewriter_config.function_optimization = off
+  rewriter_config.implementation_selector = off
+  rewriter_config.layout_optimizer = off
+  rewriter_config.loop_optimization = off
+  rewriter_config.memory_optimization = (
+      rewriter_config_pb2.RewriterConfig.NO_MEM_OPT)
+  rewriter_config.min_graph_nodes = -1
+  rewriter_config.pin_to_host_optimization = off
+  rewriter_config.remapping = off
+  rewriter_config.scoped_allocator_optimization = off
+  rewriter_config.shape_optimization = off
+
+
+def version_tuple_to_string(ver_tuple):
+  assert isinstance(ver_tuple, tuple)
+  assert len(ver_tuple) == 3
+
+  ver_tuple = [str(x) for x in ver_tuple]
+  return ".".join(ver_tuple)
+
+
+def _is_tensorrt_version_greater_equal(trt_ver, target_ver):
+  trt_ver = version.Version(version_tuple_to_string(trt_ver))
+  target_ver = version.Version(version_tuple_to_string(target_ver))
+
+  return trt_ver >= target_ver
+
+
+def is_linked_tensorrt_version_greater_equal(major, minor=0, patch=0):
+  ver = _pywrap_py_utils.get_linked_tensorrt_version()
+  return _is_tensorrt_version_greater_equal(ver, (major, minor, patch))
+
+
+def is_loaded_tensorrt_version_greater_equal(major, minor=0, patch=0):
+  ver = _pywrap_py_utils.get_loaded_tensorrt_version()
+  return _is_tensorrt_version_greater_equal(ver, (major, minor, patch))
+
+
+def is_experimental_feature_activated(feature_name):
+  """Determines if a TF-TRT experimental feature is enabled.
+
+  This helper function checks if an experimental feature was enabled using
+  the environment variable `TF_TRT_EXPERIMENTAL_FEATURES=feature_1,feature_2`.
+
+  Args:
+    feature_name: Name of the feature being tested for activation.
+  """
+
+  return (feature_name
+          in os.environ.get("TF_TRT_EXPERIMENTAL_FEATURES",
+                            default="").split(","))
+
+
+def _convert_dtype_id_to_str(dtype):
+  """Helper function to convert a dtype id to a corresponding string name."""
+  if isinstance(dtype, int):
+    return dtypes._TYPE_TO_STRING[dtype]
+  else:
+    return [dtypes._TYPE_TO_STRING[d] for d in dtype]
+
+
+def get_node_compute_dtype(node):
+  """Returns the compute DType of a GraphDef Node."""
+  # Note: Order is important, by default TF Node compute dtype is mentioned
+  # under `T` key, unless these nodes are one of ["TRTEngineOP", "Cast", "Plh"].
+  for type_key in [
+      "precision_mode",  # TRTEngineOp
+      "DstT",  # Cast Nodes
+      "dtype",  # Placeholder
+      "T",  # Everything Else
+  ]:
+    try:
+      precision_val = node.attr[type_key]
+      if type_key == "precision_mode":
+        precision_val = precision_val.s.decode("utf-8")
+        if precision_val == "":
+          continue
+        if precision_val == "FP32":
+          return "float32"
+        elif precision_val == "FP16":
+          return "float16"
+        elif precision_val == "INT8":
+          return "int8"
+        else:
+          return "unknown"
+      else:
+        return _convert_dtype_id_to_str(precision_val.type)
+    except Exception as e:
+      continue
+
+
+def get_node_io_shapes(node, key):
+  """Returns the input/output shapes of a GraphDef Node."""
+  out_shape = []
+  for shape in node.attr[key].list.shape:
+    out_shape.append([dim.size for dim in shape.dim])
+  return out_shape
+
+
+def get_trtengineop_io_dtypes(node, key):
+  """Returns the input/output dtypes of a TRTEngineOp."""
+  return _convert_dtype_id_to_str(node.attr[key].list.type)
+
+
+def get_trtengineop_io_nodes_count(node, key):
+  """Returns the number of input/output nodes of a TRTEngineOp."""
+  return len(node.attr[key].list.type)
+
+
+def get_trtengineop_node_op_count(graphdef, node_name):
+  """Counts the number of nodes and OP types of a given TRTEngineOp."""
+  ops_in_engine = collections.defaultdict(int)
+  for func in graphdef.library.function:
+    if f"{node_name}_native_segment" == func.signature.name:
+      node_count = len(func.node_def)
+      for node in func.node_def:
+        ops_in_engine[node.op] += 1
+      break
+  return node_count, ops_in_engine
+
+
+class DTypeIndex(dict):
+  """Helper class to create an index of dtypes with incremental values."""
+
+  def get_dtype_index(self, dtype):
+    if dtype not in self:
+      self[dtype] = len(self) + 1
+    return self[dtype]
+
+
+def draw_graphdef_as_graphviz(graphdef, dot_output_filename):
+  """Exports a GraphDef to GraphViz format.
+
+  - Step 1: Drawing Each Node of the compute GraphDef.
+  - Step 2: Create nodes for each collected dtype in the graph.
+  - Step 3: Creating invisible links to align properly the legend.
+
+  Each node consequently mentions:
+  - Op Type
+  - Compute Dtype
+  - Compute Device
+  """
+
+  dtype_index = DTypeIndex()
+
+  with open(dot_output_filename, "w") as f:
+    print("digraph tftrt_converted_graph {", file=f)
+
+    print("  graph [fontsize=10 fontname=\"Verdana\"];", file=f)
+    # ColorScheme Documentation: https://graphviz.org/doc/info/colors.html
+    print(
+        "  node [style=filled height=0.55 colorscheme=set312 shape=box];",
+        file=f)
+
+    # Step 1: Parsing the graph and drawing OPs one by one.
+    print("\n  subgraph tensorflow_graph {", file=f)
+    print("    node [width=1.35];", file=f)
+    nodes_with_no_inputs = []
+    for node in graphdef.node:
+      output_name = node.name
+
+      node_precision = get_node_compute_dtype(node)
+      color_idx = dtype_index.get_dtype_index(node_precision)
+
+      device_key = node.device.split("/")[-1]
+      if not device_key:
+        device_key = "device:Unspecified"
+
+      if node.op == "TRTEngineOp":
+        node_count, _ = get_trtengineop_node_op_count(graphdef, output_name)
+        node_label = f"{output_name} [{node_count}]"
+      else:
+        node_label = f"{node.op}"
+
+      # Note: double space before <br/> is necessary for formatting.
+      node_label = f"<b>{node_label}</b>  <br/><i>{device_key}</i>"
+
+      print(
+          f"    \"{output_name}\" [label=<{node_label}> "
+          f"fillcolor={color_idx}];",
+          file=f)
+
+      if len(node.input):
+        for input_full_name in node.input:
+          parts = input_full_name.split(":")
+          input_name = re.sub(r"^\^", "", parts[0])
+          print(f"  \"{input_name}\" -> \"{output_name}\";", file=f)
+      else:
+        nodes_with_no_inputs.append(output_name)
+    print("  }", file=f)
+
+    # Step 2: Creating the DType Nodes previously found in Step 1.
+    print("\n  subgraph cluster_legend {", file=f)
+    print("    label=\"Compute Dtype Legend\";", file=f)
+    print("    margin=\"30\";", file=f)
+    print("    node [width=2];", file=f)
+
+    for dtype, color_idx in dtype_index.items():
+      print(
+          f"    {dtype} [fillcolor={color_idx} label=<<b>{dtype}</b>>];",
+          file=f)
+
+    print("  }", file=f)
+
+    # Step 3: Alignment of the legend with the graph.
+    print("\n  edge[style=\"invisible\", dir=\"none\"];", file=f)
+    for dtype in dtype_index.keys():
+      for node_name in nodes_with_no_inputs:
+        print(f"  \"{dtype}\" -> \"{node_name}\"", file=f)
+
+    print("}", file=f)
+
+  print("\n===================================================================")
+  print(f"Graph Visualization Exported to: `{dot_output_filename}`.")
+  print("We recommend using https://edotor.net/ to visualize the .dot file.")
+  print("You can also use `graphviz` utility to convert them to PNG format:")
+  print("  - `sudo apt install -y graphviz`")
+  print("  - `dot -Tpng <input_filename>.dot -o <output_filename>.png`")
+  print("===================================================================\n")

tensorflow/python/compiler/xla/BUILD

@@ -0,0 +1,204 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "compiler_py",
+    srcs = [
+        "__init__.py",
+        "jit.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":xla",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "jit_test",
+    size = "small",
+    srcs = ["tests/jit_test.py"],
+    tags = [
+        "no_windows",  # TODO(b/171385770)
+    ],
+    xla_enabled = True,
+    deps = [
+        ":compiler_py",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:function",
+        "//tensorflow/python/framework:op_def_registry",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:gradients",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "xla",
+    srcs = ["xla.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/compiler/jit:xla_ops_py",
+        "//tensorflow/compiler/jit/ops:xla_ops_grad",
+        "//tensorflow/core:protos_all_py",
+        # Do not remove: required to run xla ops on Cloud.
+        "//tensorflow/compiler/tf2xla/python:xla",  # build_cleaner: keep
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/distribute:summary_op_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "xla_test",
+    srcs = ["tests/xla_test.py"],
+    tags = [
+        "no_mac",
+        "no_windows",
+    ],
+    xla_enabled = True,
+    deps = [
+        ":xla",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:control_flow_util",
+        "//tensorflow/python/ops:logging_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/summary:summary_py",
+        "//tensorflow/python/tpu:tpu_feed",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "jit_compile_test",
+    srcs = ["tests/jit_compile_test.py"],
+    tags = [
+        "no_mac",
+        "no_windows",
+    ],
+    xla_enabled = True,
+    deps = [
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "pjrt_compile_test",
+    srcs = ["tests/pjrt_compile_test.py"],
+    env = {
+        "TF_XLA_FLAGS": "--tf_xla_use_device_api --tf_xla_enable_xla_devices --tf_xla_enable_device_api_for_gpu",
+    },
+    tags = [
+        "config-cuda-only",
+        "gpu",
+        "no_oss",
+        "requires-gpu-nvidia",
+        "xla",
+    ],
+    xla_enable_strict_auto_jit = False,
+    xla_enabled = True,
+    deps = [
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+    ],
+)
+
+tf_py_strict_test(
+    name = "pjrt_autoclustering_test",
+    srcs = ["tests/pjrt_autoclustering_test.py"],
+    env = {
+        "TF_XLA_FLAGS": "--tf_xla_use_device_api --tf_xla_enable_xla_devices --tf_xla_enable_device_api_for_gpu --tf_xla_auto_jit=2 --tf_xla_enable_lazy_compilation=false --tf_xla_min_cluster_size=0",
+    },
+    tags = [
+        "config-cuda-only",
+        "gpu",
+        "no_oss",
+        "requires-gpu-nvidia",
+        "xla",
+    ],
+    xla_enable_strict_auto_jit = False,
+    xla_enabled = True,
+    deps = [
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+    ],
+)
+
+tf_py_strict_test(
+    name = "pjrt_compile_virtual_device_test",
+    srcs = ["pjrt_compile_virtual_device_test.py"],
+    env = {
+        "TF_XLA_FLAGS": "--tf_xla_use_device_api --tf_xla_enable_xla_devices --tf_xla_enable_device_api_for_gpu",
+    },
+    tags = [
+        "config-cuda-only",
+        "gpu",
+        "no_oss",
+        "requires-gpu-nvidia",
+        "xla",
+    ],
+    xla_enable_strict_auto_jit = False,
+    xla_enabled = True,
+    deps = [
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+    ],
+)

tensorflow/python/compiler/xla/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A module for controlling the Tensorflow/XLA JIT compiler."""
+
+# pylint: disable=unused-import
+from tensorflow.python.compiler.xla import jit
+from tensorflow.python.compiler.xla import xla
+# pylint: enable=unused-import

tensorflow/python/compiler/xla/experimental/BUILD

@@ -0,0 +1,65 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "xla_sharding",
+    srcs = ["xla_sharding.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/tf2xla/python:xla",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//third_party/py/numpy",
+        "@xla//xla:xla_data_proto_py",
+    ],
+)
+
+py_test(
+    name = "xla_sharding_test",
+    srcs = ["xla_sharding_test.py"],
+    strict_deps = True,
+    deps = [
+        ":xla_sharding",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//third_party/py/numpy",
+        "@xla//xla:xla_data_proto_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:variables",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "resource_variable_xla_sharding_test",
+    srcs = ["resource_variable_xla_sharding_test.py"],
+    disable_v3_4chips = False,
+    tags = ["requires-net:external"],
+    deps = [
+        ":xla_sharding",
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/tpu:device_assignment",
+        "//tensorflow/python/tpu:tpu_py",
+        "//tensorflow/python/training:adagrad",
+    ],
+)

tensorflow/python/compiler/xla/experimental/resource_variable_xla_sharding_test.py

@@ -0,0 +1,186 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.tpu import device_assignment
+from tensorflow.python.tpu import tpu
+from tensorflow.python.training import adagrad
+
+
+# Gets all the nodes of `op` in graph that have `input_node_name` as one of the
+# inputs
+def _get_op_nodes_with_input(input_node_name, op, graph):
+  nodes_with_input = []
+  for node in graph.node:
+    nodes_with_input += [
+        node
+        for input in node.input
+        if input == input_node_name and node.op == op
+    ]
+  return nodes_with_input
+
+
+# Gets XlaSharding ops connected to ReadVariableOp for the given variable_name
+def _get_xla_sharding_nodes_for_variable(variable_name, graph):
+  read_variable_op_nodes = _get_op_nodes_with_input(
+      variable_name, 'ReadVariableOp', graph
+  )
+  xla_sharding_op_nodes = []
+  for read_variable_op_node in read_variable_op_nodes:
+    xla_sharding_op_nodes += _get_op_nodes_with_input(
+        read_variable_op_node.name, 'XlaSharding', graph
+    )
+  return xla_sharding_op_nodes
+
+
+def _get_xla_sharding_proto_from_node(node):
+  sharding_proto = xla_sharding.xla_data_pb2.OpSharding()
+  sharding_proto.ParseFromString(node.attr['sharding'].s)
+  return sharding_proto
+
+
+class ResourceVariableXlaShardingTest(test.TestCase):
+
+  def setUp(self) -> None:
+    super().setUp()
+
+    context.enable_xla_sharding_for_resource_variables()
+    self.topology = tpu_cluster_resolver.initialize_tpu_system()
+    if len(config.list_logical_devices('TPU')) != 8:
+      self.skipTest('All tests require 8 TPUs.')
+
+    self.da = device_assignment.DeviceAssignment.build(
+        self.topology, computation_shape=[2, 2, 1, 2], num_replicas=1
+    )
+
+  def test_xla_sharding_ops_created_for_optimizer_slot_variables(self):
+    w = variables.Variable(
+        initial_value=math_ops.range(8, dtype=dtypes.float32),
+        name='w',
+    )
+    self.assertIsInstance(w, resource_variable_ops.BaseResourceVariable)
+    w = xla_sharding.split(
+        w,
+        split_dimension=0,
+        num_devices=8,
+    )
+    sharding_proto = xla_sharding.xla_data_pb2.OpSharding()
+    sharding_proto.ParseFromString(xla_sharding.get_tensor_sharding(w))
+    opt = adagrad.AdagradOptimizer(1.0)
+
+    @def_function.function
+    def computation(x):
+      def tpu_fn(x):
+        y = math_ops.add(w, x)
+        loss = math_ops.reduce_sum(y)
+        opt.minimize(loss, None, [w])
+        return loss
+
+      output = tpu.replicate(tpu_fn, [[x]], device_assignment=self.da)
+      return output
+
+    inputs = array_ops.reshape(math_ops.range(16, dtype=dtypes.float32), (2, 8))
+    result = computation(inputs)
+    self.assertSequenceEqual([[176.0]], self.evaluate(result))
+    graph = computation.get_concrete_function(inputs).graph.as_graph_def()
+
+    update_op_nodes = [
+        node for node in graph.node if node.op == 'ResourceApplyAdagrad'
+    ]
+    self.assertLen(update_op_nodes, 1)
+    update_op_node = update_op_nodes[0]
+
+    var_input_name = update_op_node.input[0]
+    var_sharding_nodes = _get_xla_sharding_nodes_for_variable(
+        var_input_name, graph
+    )
+    self.assertLen(var_sharding_nodes, 1)
+    self.assertProtoEquals(
+        _get_xla_sharding_proto_from_node(var_sharding_nodes[0]), sharding_proto
+    )
+
+    slot_var_input_name = update_op_node.input[1]
+    slot_var_sharding_nodes = _get_xla_sharding_nodes_for_variable(
+        slot_var_input_name, graph
+    )
+    self.assertLen(slot_var_sharding_nodes, 1)
+    self.assertProtoEquals(
+        _get_xla_sharding_proto_from_node(slot_var_sharding_nodes[0]),
+        sharding_proto,
+    )
+
+  def test_disabling_xla_sharding_ops_temporarily(self):
+    w = variables.Variable(
+        initial_value=math_ops.range(8, dtype=dtypes.float32),
+        name='w',
+    )
+    self.assertIsInstance(w, resource_variable_ops.BaseResourceVariable)
+
+    context.enable_xla_sharding_for_resource_variables()
+    with context.temporarily_disable_xla_sharding_for_resource_variables():
+      with self.assertRaisesRegex(
+          AttributeError,
+          '.*Tensor.op is undefined when eager execution is enabled.*',
+      ):
+        xla_sharding.split(
+            w,
+            split_dimension=0,
+            num_devices=8,
+        )
+
+    # xla_sharding_for_resource_variables is enabled again. Following line
+    # doesn't throw an error.
+    xla_sharding.split(
+        w,
+        split_dimension=0,
+        num_devices=8,
+    )
+
+    context.disable_xla_sharding_for_resource_variables()
+    with context.temporarily_disable_xla_sharding_for_resource_variables():
+      with self.assertRaisesRegex(
+          AttributeError,
+          '.*Tensor.op is undefined when eager execution is enabled.*',
+      ):
+        xla_sharding.split(
+            w,
+            split_dimension=0,
+            num_devices=8,
+        )
+
+    # xla_sharding_for_resource_variables stays disabled.
+    with self.assertRaisesRegex(
+        AttributeError,
+        '.*Tensor.op is undefined when eager execution is enabled.*',
+    ):
+      xla_sharding.split(
+          w,
+          split_dimension=0,
+          num_devices=8,
+      )
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/xla/experimental/xla_sharding.py

@@ -0,0 +1,642 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ======================================
+"""Experimental support for defining XLA shardings."""
+
+import numpy as _np  # Avoids becoming a part of public Tensorflow API.
+
+from tensorflow.compiler.tf2xla.python import xla as tf2xla
+from xla import xla_data_pb2
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.ops import resource_variable_ops
+
+
+class Sharding(object):
+  """A class to support adding sharding attributes to Ops.
+
+  Use the factory constructors and then call apply_to_tensor:
+    Sharding.replicate().apply_to_tensor(tensor)
+  """
+
+  def __init__(self, proto=None):
+    """Do not use this constructor; use the factory functions below."""
+    self._proto = proto
+
+  @classmethod
+  def replicate(cls):
+    """Returns a replicated sharding attribute.
+
+    This causes an op to be computed in its entirety independently on all
+    cores in the XLA device.
+    """
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(type=xla_data_pb2.OpSharding.REPLICATED))
+
+  @classmethod
+  def manual(cls):
+    """Returns a manuall sharding attribute.
+
+    This means the op is manually partitioned by the user and XLA will not
+    change the shapes.
+    """
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(type=xla_data_pb2.OpSharding.MANUAL))
+
+  @classmethod
+  def single_device(cls, core):
+    """Returns a SingleDevice sharding attribute.
+
+    This causes an op to be computed in its entirety only on one core in
+    the XLA device.
+    Args:
+      core: The core to assign this Op to.
+    """
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(
+            type=xla_data_pb2.OpSharding.MAXIMAL,
+            tile_assignment_dimensions=[1],
+            tile_assignment_devices=[core]))
+
+  @classmethod
+  def tile(cls, tile_assignment):
+    """Returns a Tiled sharding attribute.
+
+    This causes an op to be partially computed on multiple cores in the
+    XLA device.
+
+    Args:
+      tile_assignment: An np.ndarray describing the topology of the tiling and
+        which device will compute which part of the topology.
+
+    Raises:
+      TypeError: tile_assignment was not of np.array type.
+
+    TODO(jmolloy): This concept is nefarious and is not
+    something we really want to expose to users (especially as the
+    contract for tile_assignment is very strict).
+    """
+    if not isinstance(tile_assignment, _np.ndarray):
+      raise TypeError('Tile assignment must be of type np.ndarray')
+    dims = list(tile_assignment.shape)
+    flattened_devices = tile_assignment.reshape(-1, order='C')
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(
+            type=xla_data_pb2.OpSharding.OTHER,
+            tile_assignment_dimensions=dims,
+            tile_assignment_devices=list(flattened_devices)))
+
+  @classmethod
+  def subgroup_tile(cls, tile_assignment, subgroup_modes):
+    """Returns a subgroup manual sharding attribute.
+
+    This is similar to tile(), but tile_assignment has one or more dimension
+    than the tensor, and subgroup_modes define the sharding types in the last
+    dimensions of tile_assignment.
+
+    Args:
+      tile_assignment: An np.ndarray describing the topology of the tiling and
+        which device will compute which part of the topology.
+      subgroup_modes: sharding types for the dimension more than the tensor
+        shape rank.
+
+    Raises:
+      TypeError: tile_assignment was not of np.array type or subgroup_modes
+        has unsupported sharding type.
+    """
+    if not isinstance(tile_assignment, _np.ndarray):
+      raise TypeError('SubgroupTile assignment must be of type np.ndarray')
+
+    if not isinstance(subgroup_modes, list):
+      raise TypeError('subgroup_modes in subgroup manual must be of type list')
+
+    if len(tile_assignment.shape) < len(subgroup_modes):
+      raise TypeError('SubgroupTile assignment must have rank larger than'
+                      ' length of subgroup_modes')
+
+    for sharding_type in subgroup_modes:
+      if sharding_type not in [
+          xla_data_pb2.OpSharding.REPLICATED, xla_data_pb2.OpSharding.MANUAL
+      ]:
+        raise TypeError(
+            'Each sharding_type in subgroup_modes in subgroup manual must '
+            'be of type xla_data_pb2.OpSharding.REPLICATED'
+            ' or xla_data_pb2.OpSharding.MANUAL')
+    dims = list(tile_assignment.shape)
+    flattened_devices = tile_assignment.reshape(-1, order='C')
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(
+            type=xla_data_pb2.OpSharding.OTHER,
+            tile_assignment_dimensions=dims,
+            tile_assignment_devices=list(flattened_devices),
+            last_tile_dims=list(subgroup_modes)))
+
+  @classmethod
+  def partial_tile(cls, tile_assignment):
+    """Returns a partially tiled sharding attribute.
+
+    This is similar to tile(), but tile_assignment has one more dimension than
+    the tensor, and tiles in the last dimension of tile_assignment are
+    replicated.
+
+    Args:
+      tile_assignment: An np.ndarray describing the topology of the tiling and
+        which device will compute which part of the topology.
+
+    Raises:
+      TypeError: tile_assignment was not of np.array type.
+    """
+    if not isinstance(tile_assignment, _np.ndarray):
+      raise TypeError('PartialTile assignment must be of type np.ndarray')
+    dims = list(tile_assignment.shape)
+    flattened_devices = tile_assignment.reshape(-1, order='C')
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(
+            type=xla_data_pb2.OpSharding.OTHER,
+            tile_assignment_dimensions=dims,
+            tile_assignment_devices=list(flattened_devices),
+            replicate_on_last_tile_dim=True))
+
+  @classmethod
+  def split(cls, tensor, split_dimension, num_devices, input_shape=None):
+    """Returns a Sharding that splits a tensor across a dimension.
+
+    This creates a Tiled attribute, similar to tile(), but easier to use for the
+    common case of tiling a tensor N ways in one dimension.
+
+    Args:
+      tensor: A tf.Tensor to split.
+      split_dimension: The dimension number to split.
+      num_devices: The number of cores to split `tensor` over.
+      input_shape: The shape of the original tensor.
+
+    Raises:
+      ValueError: The tensor to split was smaller in the split dimension than
+        the number of devices to split over.
+    """
+    if input_shape:
+      shape = input_shape
+    else:
+      shape = tensor.shape.as_list()
+    if (shape[split_dimension] is not None and
+        shape[split_dimension] < num_devices):
+      raise ValueError('Split dimension was smaller than the required number '
+                       'of splits: shape=%r, dimension=%r, num_devices=%r' %
+                       (shape, split_dimension, num_devices))
+
+    tile_assignment_dims = [1] * len(shape)
+    tile_assignment_dims[split_dimension] = num_devices
+
+    return Sharding(
+        proto=xla_data_pb2.OpSharding(
+            type=xla_data_pb2.OpSharding.OTHER,
+            tile_assignment_dimensions=tile_assignment_dims,
+            tile_assignment_devices=range(num_devices)))
+
+  def apply_to_tensor(
+      self,
+      tensor,
+      assign_tuple_sharding=False,
+      use_sharding_op=False,
+      unspecified_dims=None,
+      sharding_v2_proto=None,
+  ):
+    """Applies this Sharding attribute to `tensor`.
+
+    Args:
+      tensor: A tf.Tensor to split.
+      assign_tuple_sharding: If the sharding type should be a tuple.
+      use_sharding_op: Whether to create a sharding op on `tensor`.
+      unspecified_dims: An optional list of dimensions unspecified.
+      sharding_v2_proto: The v2 sharding proto to use.
+
+    Returns:
+      The tensor with Sharding attribute.
+    """
+    if unspecified_dims:
+      assert use_sharding_op and not assign_tuple_sharding
+    proto = self._proto
+
+    # If passed a tf.BaseResourceVariable instead of a tf.Tensor, simply store
+    # the sharding proto on the tf.BaseResourceVariable object. An XlaShardingOp
+    # will be created down the line whenever a ReadVariableOp is created by the
+    # tf.BaseResourceVariable.
+    if (
+        isinstance(tensor, resource_variable_ops.BaseResourceVariable)
+        and context.xla_sharding_for_resource_variables_enabled()
+    ):
+      if assign_tuple_sharding:
+        proto = self._create_tuple_proto(num_outputs=1)
+      # pylint: disable=protected-access
+      tensor._set_xla_sharding(proto)
+      return tensor
+
+    if use_sharding_op:
+      if assign_tuple_sharding:
+        proto = self._create_tuple_proto(num_outputs=1)
+        tensor = tf2xla.sharding(tensor, sharding=proto.SerializeToString())
+      else:
+        tensor = tf2xla.sharding(
+            tensor,
+            sharding=proto.SerializeToString(),
+            unspecified_dims=unspecified_dims or [],
+        )
+    elif assign_tuple_sharding or len(tensor.op.outputs) > 1:
+      proto = self._get_or_create_tuple_proto(tensor.op)
+      # We can't mutate an element of old_proto.tuple_shardings, so create
+      # a new proto.
+      tuple_shardings = list(proto.tuple_shardings)
+      tuple_shardings[tensor.value_index] = self._proto
+      proto = xla_data_pb2.OpSharding(
+          type=xla_data_pb2.OpSharding.TUPLE, tuple_shardings=tuple_shardings)
+
+    # TODO(jmolloy): This need to be seriously revisited before declaring this
+    # API available for public use.
+    # pylint: disable=protected-access
+    tensor.op._set_attr('_XlaSharding',
+                        attr_value_pb2.AttrValue(s=proto.SerializeToString()))
+    if sharding_v2_proto:
+      tensor.op._set_attr(
+          '_XlaShardingV2',
+          attr_value_pb2.AttrValue(s=sharding_v2_proto.SerializeToString()),
+      )
+    return tensor
+
+  def apply_to_operation(self, operation):
+    """Applies this Sharding attribute to `operation`.
+
+    Args:
+      operation: A tf.Operation to add sharding annotation.
+    """
+    attr_value = attr_value_pb2.AttrValue(s=self._proto.SerializeToString())
+    # pylint: disable=protected-access
+    operation._set_attr('_XlaSharding', attr_value)
+
+  @property
+  def proto(self):
+    """Return the sharding protobuf of type xla_data_pb2.OpSharding."""
+    return self._proto
+
+  def _get_or_create_tuple_proto(self, op):
+    try:
+      attr = op.get_attr('_XlaSharding')
+      proto = xla_data_pb2.OpSharding()
+      proto.ParseFromString(attr)
+      return proto
+    except ValueError:
+      return self._create_tuple_proto(len(op.outputs))
+
+  def _create_tuple_proto(self, num_outputs):
+    shardings = [
+        xla_data_pb2.OpSharding(type=xla_data_pb2.OpSharding.REPLICATED)
+    ] * num_outputs
+    return xla_data_pb2.OpSharding(
+        type=xla_data_pb2.OpSharding.TUPLE, tuple_shardings=shardings)
+
+
+def copy_sharding(from_tensor, to_tensor, use_sharding_op=False):
+  """Copies the a tensor's sharding to another.
+
+  Args:
+    from_tensor: Source tensor. Must be the sole output of an op.
+    to_tensor: the tensor the annotate with the copy.
+    use_sharding_op: whether to create a sharding op on `to_tensor`.
+
+  Returns:
+    A tensor with sharding annotation copied from `from_tensor`.
+  """
+  sharding = get_tensor_sharding(from_tensor)
+  if sharding is None:
+    return to_tensor
+
+  # If passed a tf.BaseResourceVariable instead of a tf.Tensor, simply store the
+  # sharding proto on the tf.BaseResourceVariable object. An XlaShardingOp
+  # will be created down the line whenever a ReadVariableOp is created by the
+  # tf.BaseResourceVariable.
+  if (
+      isinstance(to_tensor, resource_variable_ops.BaseResourceVariable)
+      and context.xla_sharding_for_resource_variables_enabled()
+  ):
+    proto = xla_data_pb2.OpSharding()
+    proto.ParseFromString(sharding)
+    # pylint: disable=protected-access
+    to_tensor._set_xla_sharding(proto)
+    return to_tensor
+
+  if use_sharding_op:
+    to_tensor = tf2xla.sharding(to_tensor, sharding=sharding)
+  attr_value = attr_value_pb2.AttrValue(s=sharding)
+  # pylint: disable=protected-access
+  to_tensor.op._set_attr('_XlaSharding', attr_value)
+  return to_tensor
+
+# Helpers for the above factory functions that allow easy application of
+# shardings, for example:
+#   tensor = xla_sharding.replicate(tensor)
+
+
+def replicate(tensor, assign_tuple_sharding=False, use_sharding_op=False):
+  return Sharding.replicate().apply_to_tensor(
+      tensor,
+      assign_tuple_sharding=assign_tuple_sharding,
+      use_sharding_op=use_sharding_op)
+
+
+def single_device(tensor,
+                  device,
+                  assign_tuple_sharding=False,
+                  use_sharding_op=False):
+  """Returns a tensor that has SingleDevice sharding attribute."""
+  return Sharding.single_device(device).apply_to_tensor(
+      tensor,
+      assign_tuple_sharding=assign_tuple_sharding,
+      use_sharding_op=use_sharding_op)
+
+
+def tile(tensor,
+         tile_assignment,
+         assign_tuple_sharding=False,
+         use_sharding_op=False,
+         unspecified_dims=None):
+  """Returns a tensor that has tiled sharding.
+
+  Args:
+    tensor: A tf.Tensor to shard.
+    tile_assignment: An np.ndarray describing the topology of the tiling and
+      which device will compute which part of the topology.
+    assign_tuple_sharding: If the sharding type should be a tuple.
+    use_sharding_op: If true, adds a sharding op to set the sharding.
+    unspecified_dims: An optional list of dimensions unspecified.
+  """
+  return Sharding.tile(tile_assignment).apply_to_tensor(
+      tensor,
+      assign_tuple_sharding=assign_tuple_sharding,
+      use_sharding_op=use_sharding_op,
+      unspecified_dims=unspecified_dims or [])
+
+
+def split(tensor,
+          split_dimension,
+          num_devices,
+          assign_tuple_sharding=False,
+          use_sharding_op=False,
+          input_shape=None):
+  """Returns a tensor that is split along the given dimension.
+
+  Args:
+    tensor: A tf.Tensor to split.
+    split_dimension: The dimension to split.
+    num_devices: The number of devices to partition the dimension.
+    assign_tuple_sharding: If the sharding type should be a tuple.
+    use_sharding_op: If true, adds a sharding op to set the sharding.
+    input_shape: The full shape of the input tensor.
+  """
+  return Sharding.split(tensor, split_dimension, num_devices,
+                        input_shape).apply_to_tensor(
+                            tensor,
+                            assign_tuple_sharding=assign_tuple_sharding,
+                            use_sharding_op=use_sharding_op)
+
+
+def partial_tile(tensor,
+                 tile_assignment,
+                 use_sharding_op=False,
+                 unspecified_dims=None):
+  """Returns a tensor that has tiled sharding.
+
+  Args:
+    tensor: A tf.Tensor to shard.
+    tile_assignment: An np.ndarray describing the topology of the tiling and
+      which device will compute which part of the topology. It must have one
+      more dimension than tensor, and the last dimension represents partially
+      replicated tiles.
+    use_sharding_op: If true, adds a sharding op to set the sharding.
+    unspecified_dims: An optional list of dimensions unspecified.
+  """
+  return Sharding.partial_tile(tile_assignment).apply_to_tensor(
+      tensor,
+      use_sharding_op=use_sharding_op,
+      unspecified_dims=unspecified_dims or [])
+
+
+def get_op_sharding(op):
+  """Returns sharding attribute of an op.
+
+  Args:
+    op: a TensorFlow op.
+
+  Returns:
+    The attribute representing XLA sharding on this op.
+  """
+  try:
+    return op.get_attr('_XlaSharding')
+  except ValueError:
+    return None
+  except AttributeError:
+    # AttributeError: 'DistributedVarOp' object has no attribute 'get_attr'.
+    return None
+
+
+def get_tensor_sharding(tensor):
+  """Returns sharding attribute of a Tensor.
+
+  Args:
+    tensor: a Tensor.
+
+  Returns:
+    The attribute representing XLA sharding on tensor's op.
+  """
+  # If passed a tf.BaseResourceVariable instead of a tf.Tensor, simply get the
+  # sharding proto set on the _xla_sharding field of the tf.BaseResourceVariable
+  # object.
+  if (
+      isinstance(tensor, resource_variable_ops.BaseResourceVariable)
+      and context.xla_sharding_for_resource_variables_enabled()
+  ):
+    # pylint: disable=protected-access
+    sharding = tensor._get_xla_sharding()
+    if sharding is None:
+      return None
+    else:
+      return sharding.SerializeToString()
+
+  try:
+    return get_op_sharding(tensor.op)
+  except AttributeError:
+    # AttributeError: Tensor.op is meaningless when eager execution is enabled.
+    return None
+
+
+def get_sharding_tile_shape(sharding):
+  """Returns the tile assignment shape for a sharded Tensor.
+
+  Args:
+    sharding: a serialized OpSharding message describing the layout of a
+      sharded Tensor.
+
+  Returns:
+    A list, for each dimension of the sharded Tensor, of the number of shards
+      into which it has been split. Returns None if the input indicates no tile
+      assignments.
+  """
+  if sharding is None:
+    return None
+  sharding_message = xla_data_pb2.OpSharding()
+  sharding_message.ParseFromString(sharding)
+  if sharding_message.tile_assignment_dimensions:
+    return sharding_message.tile_assignment_dimensions
+  else:
+    return None
+
+
+def auto_to_manual_spmd_partition(tensor,
+                                  manual_sharding,
+                                  single_dim=-1,
+                                  unspecified_dims=None):
+  """Switches from automatic SPMD partitioning to manual partitioning.
+
+  Converts a full-shaped tensor (to be automatically partitioned by SPMD
+  partitioner) to a shard-shaped tensor to be consumed by manually partitioned
+  ops.
+
+  Args:
+    tensor: A tf.Tensor in full shape.
+    manual_sharding: A serialized string of OpSharding to be used in manual
+      partitioning.
+    single_dim: If >= 0, the conversion will happen only on this dim in
+      subgroups.
+    unspecified_dims: An optional list of dimensions unspecified.
+
+  Returns:
+    A shard-shaped tensor to be consumed by manually partitioned ops.
+  """
+  return tf2xla.spmd_full_to_shard_shape(
+      tensor,
+      manual_sharding=manual_sharding,
+      dim=single_dim,
+      unspecified_dims=unspecified_dims or [])
+
+
+def manual_to_auto_spmd_partition(tensor,
+                                  manual_sharding,
+                                  full_shape,
+                                  single_dim=-1,
+                                  unspecified_dims=None):
+  """Switches from manual partitioning to automatic SPMD partitioning.
+
+  Converts a shard-shaped tensor (manually partitioned in SPMD-style) to a
+  full-shaped tensor to be partitioned automatically by the SPMD partitioner.
+
+  Args:
+    tensor: A tf.Tensor in shard shape.
+    manual_sharding: a serialized string of OpSharding to be used in manual
+      partitioning.
+    full_shape: the shape of tensor before partitioning.
+    single_dim: If >= 0, the conversion will happen only on this dim in
+      subgroups.
+    unspecified_dims: An optional list of dimensions unspecified.
+
+  Returns:
+    A full-shaped tensor to be partitioned automatically by the SPMD
+    partitioner.
+  """
+  return tf2xla.spmd_shard_to_full_shape(
+      tensor,
+      manual_sharding=manual_sharding,
+      full_shape=full_shape,
+      dim=single_dim,
+      unspecified_dims=unspecified_dims or [])
+
+
+def mesh_split_sharding(device_mesh,
+                        tensor_split_dims_mapping,
+                        manual_mesh_dims=None):
+  """Returns a Sharding object representing sharding along multiple dimensions.
+
+  Args:
+    device_mesh: An np.ndarray describing the topology of the device mesh and
+      each element is the ID of the device in the topology.
+    tensor_split_dims_mapping: A list of integers that map each tensor axis to
+      the device mesh axis along which it is sharded. Its length is the tensor
+      rank, and tensor_split_dims_mapping[i] is device mesh axis for tensor
+      dimension i. Use -1 for tensor dimensions that are not sharded.
+    manual_mesh_dims: An optional list of mesh dims for manual subgroups.
+
+  Raises:
+    ValueError: The number of tensor split dimensions is larger than device mesh
+      rank.
+  """
+  manual_mesh_dims = manual_mesh_dims or []
+  permutation = [d for d in tensor_split_dims_mapping if d >= 0
+                ] + manual_mesh_dims
+  if len(permutation) > len(device_mesh.shape):
+    raise ValueError(
+        'Number of tensor split dimensions (%r) is larger than device mesh '
+        'rank (%r). tensor_split_dims_mapping: %r, device_mesh.shape: %r' %
+        (len(permutation), len(
+            device_mesh.shape), tensor_split_dims_mapping, device_mesh.shape))
+  # Append replicated dimensions to the end.
+  transpose_permutation = permutation + [
+      d for d in range(len(device_mesh.shape)) if d not in permutation
+  ]
+  tile_assignment = _np.transpose(device_mesh, transpose_permutation)
+  tile_shape = [
+      1 if d < 0 else device_mesh.shape[d]
+      for d in (tensor_split_dims_mapping + manual_mesh_dims)
+  ]
+  subgroup_modes = [xla_data_pb2.OpSharding.MANUAL] * len(manual_mesh_dims)
+  partial = len(permutation) < len(device_mesh.shape)
+  if partial:
+    tile_shape.append(_np.prod(device_mesh.shape) // _np.prod(tile_shape))
+    subgroup_modes.append(xla_data_pb2.OpSharding.REPLICATED)
+  tile_assignment = _np.reshape(tile_assignment, tile_shape)
+
+  if manual_mesh_dims:
+    return Sharding.subgroup_tile(tile_assignment, subgroup_modes)
+
+  if partial:
+    return Sharding.partial_tile(tile_assignment)
+  return Sharding.tile(tile_assignment)
+
+
+def mesh_split(tensor,
+               device_mesh,
+               tensor_split_dims_mapping,
+               use_sharding_op=False,
+               manual_mesh_dims=None,
+               unspecified_dims=None):
+  """Returns a tensor that is split along multiple dimensions in a device mesh.
+
+  Args:
+    tensor: A tf.Tensor to split.
+    device_mesh: An np.ndarray describing the topology of the device mesh and
+      each element is the ID of the device in the topology.
+    tensor_split_dims_mapping: A list of integers that map each tensor axis to
+      the device mesh axis along which it is sharded. Its length is the tensor
+      rank, and tensor_split_dims_mapping[i] is device mesh axis for tensor
+      dimension i. Use -1 for tensor dimensions that are not sharded.
+    use_sharding_op: If true, adds a sharding op to set the sharding.
+    manual_mesh_dims: An optional list of mesh dims for manual subgroups.
+    unspecified_dims: An optional list of dimensions unspecified.
+
+  Raises:
+    ValueError: The number of tensor split dimensions is larger than device mesh
+      rank.
+  """
+  sharding = mesh_split_sharding(device_mesh, tensor_split_dims_mapping,
+                                 manual_mesh_dims)
+  return sharding.apply_to_tensor(
+      tensor,
+      use_sharding_op=use_sharding_op,
+      unspecified_dims=unspecified_dims or [])

tensorflow/python/compiler/xla/experimental/xla_sharding_test.py

@@ -0,0 +1,200 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ======================================
+"""Tests for xla_sharding.Sharding class and associated module functions."""
+
+from absl.testing import absltest
+import numpy as np
+
+from google.protobuf.message import DecodeError
+from xla import xla_data_pb2
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+
+
+class ShardingTest(test_util.TensorFlowTestCase):
+  """Tests for member functions of the class xla_sharding.Sharding."""
+
+  def test_sharding_is_default_constructable(self):
+    sharding = xla_sharding.Sharding()
+    self.assertIsNotNone(sharding)
+
+  def test_sharding_factory_functions_can_return_sharding_objects(self):
+    """Tests the various recommended ways to construct a Sharding object.
+
+    This is the most minimal of tests, doesn't assert anything about the
+    Sharding object produced by a given factory methods other than that it
+    has the correct type.
+    """
+    self.assertIsInstance(xla_sharding.Sharding.replicate(),
+                          xla_sharding.Sharding)
+    self.assertIsInstance(xla_sharding.Sharding.manual(), xla_sharding.Sharding)
+    self.assertIsInstance(
+        xla_sharding.Sharding.single_device(0), xla_sharding.Sharding)
+    self.assertIsInstance(
+        xla_sharding.Sharding.tile(np.ones([3], dtype=int)),
+        xla_sharding.Sharding)
+    self.assertIsInstance(
+        xla_sharding.Sharding.partial_tile(np.ones([3], dtype=int)),
+        xla_sharding.Sharding)
+    self.assertIsInstance(
+        xla_sharding.Sharding.split(
+            array_ops.ones([3, 8, 7], dtype=dtypes.int32), 1, 2),
+        xla_sharding.Sharding)
+    self.assertIsInstance(
+        xla_sharding.Sharding.subgroup_tile(
+            np.ones([2, 3, 3], dtype=int), [
+                xla_data_pb2.OpSharding.REPLICATED,
+                xla_data_pb2.OpSharding.MANUAL
+            ]), xla_sharding.Sharding)
+
+
+class XlaShardingTest(test_util.TensorFlowTestCase):
+  """Tests for non-member functions in the module xla_sharding.py."""
+
+  def setUp(self):
+    super().setUp()
+    context.enable_xla_sharding_for_resource_variables()
+
+  def _graph_has_xla_sharding_op(self, graph):
+    for node in graph.node:
+      if node.op == 'XlaSharding' and any(
+          'ReadVariableOp' in input for input in node.input
+      ):
+        return True
+
+    return False
+
+  def test_replicate_annotates_tensor_correctly(self):
+
+    @def_function.function
+    def replicate_helper(tensor):
+      self.assertIsNone(xla_sharding.get_tensor_sharding(tensor))
+      replicated_tensor = xla_sharding.replicate(tensor)
+      replicated_sharding = xla_sharding.get_tensor_sharding(replicated_tensor)
+      self.assertIsNotNone(replicated_sharding)
+      self.assertIsNone(
+          xla_sharding.get_sharding_tile_shape(replicated_sharding))
+      return replicated_tensor
+
+    in_tensor = array_ops.ones([4, 5, 6], dtype=dtypes.float32)
+    result = replicate_helper(array_ops.ones([4, 5, 6], dtype=dtypes.float32))
+    self.assertAllEqual(in_tensor, result)
+
+    var = variables.Variable(initial_value=in_tensor, name='var')
+    graph = replicate_helper.get_concrete_function(var).graph.as_graph_def()
+    self.assertTrue(self._graph_has_xla_sharding_op(graph))
+
+  def test_tile_annotates_tensor_correctly(self):
+
+    @def_function.function
+    def tile_helper(tensor):
+      self.assertIsNone(xla_sharding.get_tensor_sharding(tensor))
+      tiled_tensor = xla_sharding.tile(tensor, np.array([2, 1, 6]))
+      self.assertIsInstance(tiled_tensor, type(tensor))
+      tiled_sharding = xla_sharding.get_tensor_sharding(tiled_tensor)
+      tile_shape = xla_sharding.get_sharding_tile_shape(tiled_sharding)
+      # This is the shape of the tile assignment [2, 1, 6]
+      expected_shape = [3]
+      self.assertEqual(expected_shape, tile_shape)
+      return tiled_tensor
+
+    in_tensor = array_ops.ones([4, 5, 6], dtype=dtypes.float32)
+    result = tile_helper(array_ops.ones([4, 5, 6], dtype=dtypes.float32))
+    self.assertAllEqual(in_tensor, result)
+
+    var = variables.Variable(initial_value=in_tensor, name='var')
+    graph = tile_helper.get_concrete_function(var).graph.as_graph_def()
+    self.assertTrue(self._graph_has_xla_sharding_op(graph))
+
+  def test_split_annotates_tensor_correctly(self):
+
+    @def_function.function
+    def split_helper(tensor):
+      self.assertIsNone(xla_sharding.get_tensor_sharding(tensor))
+      split_tensor = xla_sharding.split(tensor, 2, 3)
+      self.assertIsInstance(split_tensor, type(tensor))
+      split_sharding = xla_sharding.get_tensor_sharding(split_tensor)
+      split_shape = xla_sharding.get_sharding_tile_shape(split_sharding)
+      expected_shape = [1, 1, 3]
+      self.assertEqual(expected_shape, split_shape)
+      return split_tensor
+
+    in_tensor = array_ops.ones([4, 5, 6], dtype=dtypes.float32)
+    result = split_helper(array_ops.ones([4, 5, 6], dtype=dtypes.float32))
+    self.assertAllEqual(in_tensor, result)
+
+    var = variables.Variable(initial_value=in_tensor, name='var')
+    graph = split_helper.get_concrete_function(var).graph.as_graph_def()
+    self.assertTrue(self._graph_has_xla_sharding_op(graph))
+
+  def test_split_raises_error_with_incommensurate_dimensions(self):
+
+    @def_function.function
+    def split_helper(tensor):
+      split_tensor = xla_sharding.split(tensor, 0, 8)
+      return split_tensor
+
+    with self.assertRaises(ValueError):
+      _ = split_helper(array_ops.ones([4, 5, 6], dtype=dtypes.float32))
+
+    # TODO(drm): Modify split() so that this call raises an error since
+    # 8 does not divide 9 (currently only checks that 8 is smaller than 9,
+    # which it is, but this is not good for splitting).
+    # with self.assertRaises(ValueError):
+    #   _ = split_helper(array_ops.ones([9, 5, 6], dtype=dtypes.float32))
+
+  def test_copy_sharding_succeeds_with_identically_shaped_tensors(self):
+
+    @def_function.function
+    def copy_helper(tensor):
+      tensor_src = array_ops.identity(tensor)
+      tensor_src = xla_sharding.split(tensor, 2, 3)
+      sharding_src = xla_sharding.get_tensor_sharding(tensor_src)
+      shape_src = xla_sharding.get_sharding_tile_shape(sharding_src)
+      self.assertEqual([1, 1, 3], shape_src)
+
+      tensor_dest = array_ops.identity(tensor)
+      self.assertIsNone(xla_sharding.get_tensor_sharding(tensor_dest))
+
+      xla_sharding.copy_sharding(tensor_src, tensor_dest)
+      sharding_dest = xla_sharding.get_tensor_sharding(tensor_dest)
+      shape_dest = xla_sharding.get_sharding_tile_shape(sharding_dest)
+      self.assertEqual([1, 1, 3], shape_dest)
+      return tensor_dest
+
+    in_tensor = array_ops.ones([4, 5, 6], dtype=dtypes.float32)
+    result = copy_helper(array_ops.ones([4, 5, 6], dtype=dtypes.float32))
+    self.assertAllEqual(in_tensor, result)
+
+    var = variables.Variable(initial_value=in_tensor, name='var')
+    graph = copy_helper.get_concrete_function(var).graph.as_graph_def()
+    self.assertTrue(self._graph_has_xla_sharding_op(graph))
+
+  def test_get_sharding_tile_shape_returns_none_on_none_input(self):
+    self.assertIsNone(xla_sharding.get_sharding_tile_shape(None))
+
+  def test_get_sharding_tile_shape_raises_error_on_nonparsable_input(self):
+    bad_proto_data = b'\x0f'
+    with self.assertRaises(DecodeError):
+      xla_sharding.get_sharding_tile_shape(bad_proto_data)
+
+
+if __name__ == '__main__':
+  absltest.main()

tensorflow/python/compiler/xla/jit.py

@@ -0,0 +1,156 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Library for controlling the Tensorflow/XLA JIT compiler."""
+
+import contextlib
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.util.tf_export import tf_export
+
+
+_XLA_SCOPE_KEY = ("__xla_scope",)
+
+
+class _XlaScope(object):
+  """Keeps track of previous XLA scope calls, and depth of current call."""
+
+  def __init__(self, count, depth):
+    self.count = count
+    self.depth = depth
+
+
+@contextlib.contextmanager
+@tf_export("xla.experimental.jit_scope")
+def experimental_jit_scope(compile_ops=True, separate_compiled_gradients=False):
+  """Enable or disable JIT compilation of operators within the scope.
+
+  NOTE: This is an experimental feature.
+
+  The compilation is a hint and only supported on a best-effort basis.
+
+  Example usage:
+
+    ```python
+    with tf.xla.experimental.jit_scope():
+      c = tf.matmul(a, b)  # compiled
+    with tf.xla.experimental.jit_scope(compile_ops=False):
+      d = tf.matmul(a, c)  # not compiled
+    with tf.xla.experimental.jit_scope(
+        compile_ops=lambda node_def: 'matmul' in node_def.op.lower()):
+      e = tf.matmul(a, b) + d  # matmul is compiled, the addition is not.
+    ```
+
+  Example of `separate_compiled_gradients`:
+
+    ```python
+    # In the example below, the computations for f, g and h will all be compiled
+    # in separate scopes.
+    with tf.xla.experimental.jit_scope(
+        separate_compiled_gradients=True):
+      f = tf.matmul(a, b)
+    g = tf.gradients([f], [a, b], name='mygrads1')
+    h = tf.gradients([f], [a, b], name='mygrads2')
+    ```
+
+  Ops that are not in the scope may be clustered and compiled with ops in
+  the scope with `compile_ops=True`, while the ops in the scope with
+  `compile_ops=False` will never be compiled.
+
+  For example:
+
+    ```python
+    # In the example below, x and loss may be clustered and compiled together,
+    # while y will not be compiled.
+    with tf.xla.experimental.jit_scope():
+      x = tf.matmul(a, b)
+    with tf.xla.experimental.jit_scope(compile_ops=False):
+      y = tf.matmul(c, d)
+    loss = x + y
+    ```
+
+  If you want to only compile the ops in the scope with `compile_ops=True`,
+  consider adding an outer `jit_scope(compile_ops=False)`:
+
+    ```python
+    # In the example below, only x will be compiled.
+    with tf.xla.experimental.jit_scope(compile_ops=False):
+      with tf.xla.experimental.jit_scope():
+        x = tf.matmul(a, b)
+      y = tf.matmul(c, d)
+      loss = x + y
+    ```
+
+  Args:
+    compile_ops: Whether to enable or disable compilation in the scope.
+      Either a Python bool, or a callable that accepts the parameter
+      `node_def` and returns a python bool.
+    separate_compiled_gradients: If true put each gradient subgraph into a
+      separate compilation scope. This gives fine-grained control over which
+      portions of the graph will be compiled as a single unit. Compiling
+      gradients separately may yield better performance for some graphs.
+      The scope is named based on the scope of the forward computation as well
+      as the name of the gradients. As a result, the gradients will be compiled
+      in a scope that is separate from both the forward computation, and from
+      other gradients.
+  Raises:
+    RuntimeError: if called when eager execution is enabled.
+  Yields:
+    The current scope, enabling or disabling compilation.
+  """
+  if context.executing_eagerly():
+    raise RuntimeError("xla.experimental.jit_scope is not supported when eager "
+                       "execution is enabled. Try use it inside tf.function.")
+
+  if callable(compile_ops):
+    def xla_compile(node_def):
+      return attr_value_pb2.AttrValue(b=compile_ops(node_def))
+  else:
+    xla_compile = attr_value_pb2.AttrValue(b=compile_ops)
+
+  attrs = {
+      "_XlaCompile":
+          xla_compile,
+      "_XlaSeparateCompiledGradients":
+          attr_value_pb2.AttrValue(b=bool(separate_compiled_gradients))
+  }
+
+  # Find the singleton counter for the current scoped graph.  If it
+  # doesn't exist, create one.
+  xla_scope_counter = ops.get_collection(_XLA_SCOPE_KEY)
+  if not xla_scope_counter:
+    xla_scope_counter = _XlaScope(0, 0)
+    ops.add_to_collection(_XLA_SCOPE_KEY, xla_scope_counter)
+  else:
+    xla_scope_counter = xla_scope_counter[0]
+
+  if xla_scope_counter.depth == 0:
+    # If we're at the root xla scope, we can increase the counter so
+    # future calls to jit_scope use a different scope value.
+    # If we're already within a scope, we'll be fusing using the scope
+    # controlled by the parent.
+    attrs["_XlaScope"] = attr_value_pb2.AttrValue(
+        s=("jit_scope_%d" % xla_scope_counter.count).encode())
+    xla_scope_counter.count += 1
+
+  xla_scope_counter.depth += 1
+
+  # pylint: disable=protected-access
+  with ops.get_default_graph()._attr_scope(attrs):
+    yield
+  # pylint: enable=protected-access
+
+  xla_scope_counter.depth -= 1

tensorflow/python/compiler/xla/pjrt_compile_virtual_device_test.py

@@ -0,0 +1,75 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests virtual device compilation + execution using the Device API (aka PjRt).
+
+This feature is still under active development and is protected behind the
+`--tf_xla_use_device_api` flag in the `TF_XLA_FLAGS` environment variable.
+"""
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+
+
+class PjrtCompileVirtualDeviceTest(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    gpus = config.list_physical_devices("GPU")
+    config.set_logical_device_configuration(
+        gpus[0],
+        [
+            context.LogicalDeviceConfiguration(memory_limit=1024),
+            context.LogicalDeviceConfiguration(memory_limit=1024),
+            context.LogicalDeviceConfiguration(memory_limit=1024),
+        ],
+    )
+
+  def test_xla_launch_and_tf_kernel_on_gpu_device(self):
+    @def_function.function(jit_compile=True)
+    def foo(x, y):
+      return x + y + 1
+
+    @def_function.function(jit_compile=True)
+    def bar(x, y):
+      x.assign(y)
+      y.assign_add([1.0, 1.0])
+
+    with ops.device("/device:GPU:1"):
+      a = constant_op.constant([1.0, 2.0])
+      x = variables.Variable([0.0, 1.0])
+      result_tensor = foo(x, a)
+
+    self.assertAllClose(result_tensor.numpy(), [2.0, 4.0], atol=1e-05)
+
+    # The following use case is tested:
+    # Variable updates following an XLA computation that reads the updated
+    # variables.
+    with ops.device("/device:GPU:1"):
+      var_a = variables.Variable([0.0, 1.0])
+      var_b = variables.Variable([1.0, 2.0])
+      bar(var_a, var_b)
+      result = foo(var_a, var_b)
+
+    self.assertAllClose([1.0, 2.0], var_a.value(), atol=1e-05)
+    self.assertAllClose([2.0, 3.0], var_b.value(), atol=1e-05)
+    self.assertAllClose(result, [4.0, 6.0], atol=1e-05)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/jit_compile_test.py

@@ -0,0 +1,101 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.platform import test
+
+
+class JitCompileTest(test.TestCase):
+
+  def testBasic(self):
+    with ops.Graph().as_default() as g:
+
+      def fn(x, a):
+        return x + a
+
+      xla_func = def_function.function(fn, jit_compile=True)
+      inputs = array_ops.placeholder(dtypes.float32, [5])
+      x = xla_func(inputs, 1)
+      with session.Session(graph=g) as sess:
+        y = sess.run(x, feed_dict={inputs: [1, 2, 2, 3, 3]})
+        self.assertTrue(x.graph.as_graph_def().library.function[0]
+                        .attr["_XlaMustCompile"].b)
+        self.assertAllClose([2, 3, 3, 4, 4], y)
+
+  def testDerivative(self):
+    def fn(x, a):
+      return 2 * x + a
+
+    with ops.Graph().as_default() as g:
+      xla_func = def_function.function(fn, jit_compile=True)
+      with backprop.GradientTape() as tape:
+        inputs = array_ops.placeholder(dtypes.float32, [5])
+        tape.watch(inputs)
+        outputs = xla_func(inputs, 1)
+      grads = tape.gradient(outputs, inputs)
+
+    with session.Session(graph=g) as sess:
+      grads_tensor = sess.run(grads, feed_dict={inputs: [1, 2, 2, 3, 3]})
+      self.assertAllClose([2, 2, 2, 2, 2], grads_tensor)
+      (forward, backward) = xla_func.get_concrete_function(
+          inputs, 1)._delayed_rewrite_functions.forward_backward()
+
+      # Check that the must-compile attribute gets correctly propagated to the
+      # created derivatives.
+      self.assertTrue(forward.cached_definition.attr["_XlaMustCompile"])
+      self.assertTrue(backward.function_def.attr["_XlaMustCompile"])
+
+  def testBasicInt32(self):
+    with ops.Graph().as_default() as g:
+
+      def fn(x, a):
+        return x + a
+
+      xla_func = def_function.function(fn, jit_compile=True)
+      inputs = array_ops.placeholder(dtypes.int32, [5])
+      x = xla_func(inputs, 1)
+      with session.Session(graph=g) as sess:
+        y = sess.run(x, feed_dict={inputs: [1, 2, 2, 3, 3]})
+        self.assertTrue(x.graph.as_graph_def().library.function[0]
+                        .attr["_XlaMustCompile"].b)
+        self.assertAllClose([2, 3, 3, 4, 4], y)
+
+  # Checking that we crash on an unsupported operation lets us test that the XLA
+  # compiler was actually invoked.
+  def testUnsupportedOps(self):
+    with ops.Graph().as_default() as g:
+
+      def fn(x):
+        return string_ops.string_length(
+            string_ops.string_format('{}', x))
+
+      xla_func = def_function.function(fn, jit_compile=True)
+      inputs = array_ops.placeholder(dtypes.float32, [5])
+      x = xla_func(inputs)
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  "Detected unsupported operations"):
+        with session.Session(graph=g) as sess:
+          sess.run(x, feed_dict={inputs: [1, 2, 2, 3, 3]})
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/jit_test.py

@@ -0,0 +1,315 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for python.compiler.xla.jit."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.compiler.xla import jit
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import function
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gradients
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def enable_jit_nonstateful(node_def):
+  op_def = op_def_registry.get(node_def.op)
+  if op_def is None:
+    raise ValueError("Unregistered op being created: %s" % node_def)
+
+  return not op_def.is_stateful
+
+
+class JITTest(test.TestCase, parameterized.TestCase):
+
+  def compute(self, use_jit, compute_fn):
+    random_seed.set_random_seed(1234)
+    with self.session(graph=ops.Graph()) as sess:
+      with jit.experimental_jit_scope(use_jit):
+        r = compute_fn()
+      sess.run(variables.global_variables_initializer())
+      return (r, sess.run(r))
+
+  @test_util.run_v2_only
+  def testJITInEager(self):
+
+    with self.assertRaisesRegex(
+        RuntimeError, "xla.experimental.jit_scope is not supported when eager "
+        "execution is enabled. Try use it inside tf.function."):
+      with jit.experimental_jit_scope(True):
+        constant_op.constant(1)
+
+  @test_util.build_as_function_and_v1_graph
+  def testJITCreateOpsLambda(self):
+    """Test several ways of customizing the compilation attribute."""
+    def create_ops():
+      with variable_scope.variable_scope(
+          "root",
+          initializer=init_ops.random_uniform_initializer(
+              -0.1, 0.1, seed=2)):
+        inputs = random_ops.random_uniform((1,), minval=-10, maxval=10, seed=1)
+        return inputs
+    v_false_1_t, v_false_1 = self.compute(False, create_ops)
+    _, v_false_2 = self.compute(False, create_ops)
+    v_true_1_t, v_true_1 = self.compute(enable_jit_nonstateful, create_ops)
+    _, v_true_2 = self.compute(enable_jit_nonstateful, create_ops)
+    v_all_true_t, _ = self.compute(True, create_ops)
+    self.assertFalse(v_false_1_t.op.get_attr("_XlaCompile"))
+    v_true_1_t_sampler_op = v_true_1_t.graph.get_operation_by_name(
+        "root/random_uniform/RandomUniform")
+    v_all_true_t_sampler_op = v_all_true_t.graph.get_operation_by_name(
+        "root/random_uniform/RandomUniform")
+
+    self.assertFalse(v_true_1_t_sampler_op.get_attr("_XlaCompile"))
+    self.assertTrue(v_all_true_t_sampler_op.get_attr("_XlaCompile"))
+
+    self.assertTrue(v_true_1_t.op.get_attr("_XlaCompile"))
+    self.assertTrue(v_all_true_t.op.get_attr("_XlaCompile"))
+
+    # Additionally ensure that where no JIT compilation happens on the
+    # random_uniform op, the output values are identical to the case
+    # where no JIT compilation happens anywhere.
+    self.assertAllClose(v_false_1, v_false_2)
+    self.assertAllClose(v_true_1, v_true_2)
+    self.assertAllClose(v_false_1, v_true_1)
+
+  @test_util.build_as_function_and_v1_graph
+  def testJITXlaScope(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope(True):
+        # XlaScope 0
+        a1 = constant_op.constant(1)
+      with jit.experimental_jit_scope(True):
+        # XlaScope 1
+        a2 = constant_op.constant(1)
+        with jit.experimental_jit_scope(True):
+          # XlaScope still 1, depth 1
+          a3 = constant_op.constant(1)
+          with jit.experimental_jit_scope(True):
+            # XlaScope still 1, depth 2
+            a4 = constant_op.constant(1)
+          # XlaScope still 1, depth 1
+          a5 = constant_op.constant(1)
+      with jit.experimental_jit_scope(True):
+        # XlaScope now 2, depth 0
+        a6 = constant_op.constant(1)
+
+    self.assertEqual(b"jit_scope_0", a1.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a2.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a3.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a4.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a5.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_2", a6.op.get_attr("_XlaScope"))
+
+  @test_util.build_as_function_and_v1_graph
+  def testJITVariableSeed(self):
+    """Test that the stateful initializer is not marked for compilation.
+
+    XLA does not currently support seeded initialization and XLA initializers
+    therefore return different values than non-XLA counterparts.  Here
+    we ensure that if we can disable JIT compilation for the initializers and
+    get the same variable values as if no JIT compilation happened.
+    """
+    def create_ops():
+      with variable_scope.variable_scope(
+          "root",
+          initializer=init_ops.random_uniform_initializer(
+              -0.1, 0.1, seed=2)):
+        inputs = variable_scope.get_variable("var", (1,))
+        return inputs
+    _, v_false_1 = self.compute(False, create_ops)
+    _, v_false_2 = self.compute(False, create_ops)
+    _, v_true_1 = self.compute(enable_jit_nonstateful, create_ops)
+    _, v_true_2 = self.compute(enable_jit_nonstateful, create_ops)
+    self.assertAllClose(v_false_1, v_false_2)
+    self.assertAllClose(v_true_1, v_true_2)
+    self.assertAllClose(v_false_1, v_true_1)
+
+  @test_util.build_as_function_and_v1_graph
+  def testDefunNoJitScope(self):
+    with self.session(graph=ops.Graph()):
+
+      @function.Defun(compiled=True, noinline=True)
+      def mulop(x1, x2):
+        return x1 * x2
+      x = constant_op.constant(1.0)
+      r = mulop(x, x)
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      # No enclosing jit scope so function sets its own value for _XlaScope.
+      self.assertEqual(b"function_mulop", func_attrs["_XlaScope"].s)
+
+  @test_util.build_as_function_and_v1_graph
+  def testDefunInheritsJitScope(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope(True):
+        @function.Defun(compiled=True, noinline=True)
+        def mulop(x1, x2):
+          return x1 * x2
+        x = constant_op.constant(1.0)
+        r = mulop(x, x)
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      # Ensure _XlaScope is inherited from enclosing context.
+      self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s)
+
+
+class CompilationEnabledInGradientTest(test.TestCase, parameterized.TestCase):
+
+  @test_util.build_as_function_and_v1_graph
+  def testCompilationInGradient(self):
+    with self.cached_session():
+      x = constant_op.constant([[3.]])
+      y_nc = math_ops.matmul(x, x, name="not_compiled")
+      with jit.experimental_jit_scope():
+        y_c = math_ops.matmul(y_nc, y_nc, name="compiled")
+      x_grads = gradients.gradients([y_c], [x])[0]
+      operations = x.graph.get_operations()
+      c_grad_ops = [
+          op for op in operations if "gradients/compiled" in op.name]
+      nc_grad_ops = [
+          op for op in operations if "gradients/not_compiled" in op.name]
+      self.assertGreater(len(c_grad_ops), 0)
+      self.assertGreater(len(nc_grad_ops), 0)
+      for cg in c_grad_ops:
+        self.assertTrue(cg.get_attr("_XlaCompile"))
+      for ncg in nc_grad_ops:
+        with self.assertRaisesRegex(ValueError, "[Nn]o attr named"):
+          ncg.get_attr("_XlaCompile")
+
+      # d/dx (x ** 4) = 4 * (x ** 3)
+      self.assertAllClose([[108]], x_grads)
+
+  @test_util.build_as_function_and_v1_graph
+  def testCompilationGradientScopeNames(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope():
+        # XlaScope 0
+        a1 = constant_op.constant([[1.]])
+        a1t = math_ops.matmul(a1, a1)
+      with jit.experimental_jit_scope():
+        # XlaScope 1
+        a2 = constant_op.constant([[1.]])
+        a2t = math_ops.matmul(a2, a2)
+
+      self.assertEqual(b"jit_scope_0", a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1", a2.op.get_attr("_XlaScope"))
+      grad_a1 = gradients.gradients(a1t, a1, name="GA")[0]
+      grad_a2 = gradients.gradients(a2t, a2, name="GB")[0]
+      grad_a1 = grad_a1.op.inputs[0]
+      grad_a2 = grad_a2.op.inputs[0]
+      self.assertTrue(grad_a1.op.get_attr("_XlaCompile"))
+      self.assertTrue(grad_a2.op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0", grad_a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1", grad_a2.op.get_attr("_XlaScope"))
+
+  @test_util.build_as_function_and_v1_graph
+  def testCompilationSeparateGradientScopeNames(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope(True, separate_compiled_gradients=True):
+        # XlaScope 0
+        a1 = constant_op.constant([[1.]])
+        a1t = math_ops.matmul(a1, a1)
+      with jit.experimental_jit_scope(True, separate_compiled_gradients=True):
+        # XlaScope 1
+        a2 = constant_op.constant([[1.]])
+        a2t = math_ops.matmul(a2, a2)
+
+      self.assertEqual(b"jit_scope_0", a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1", a2.op.get_attr("_XlaScope"))
+      grad_a1 = gradients.gradients(a1t, a1, name="GA")[0]
+      grad_a2 = gradients.gradients(a2t, a2, name="GB")[0]
+      grad_a1 = grad_a1.op.inputs[0]
+      grad_a2 = grad_a2.op.inputs[0]
+      self.assertTrue(grad_a1.op.get_attr("_XlaCompile"))
+      self.assertTrue(grad_a2.op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0_grad_GA",
+                       grad_a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1_grad_GB",
+                       grad_a2.op.get_attr("_XlaScope"))
+
+  @test_util.build_as_function_and_v1_graph
+  def testPlaysNicelyWithDefun(self):
+    with self.session(graph=ops.Graph()) as sess:
+      with jit.experimental_jit_scope(True):
+        @function.Defun(compiled=True, noinline=True)
+        def mulop(x1, x2):
+          return x1 * x2
+        x = constant_op.constant(1.0)
+        r = mulop(x, x)
+        g_r = gradients.gradients(r, x, name="GA")[0]
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s)
+
+      # Ensure the gradient (SymbolicGradient) is compiled, with the same
+      # _XlaScope as the function itself.
+      grad_op = g_r.op.inputs[0].op
+      self.assertTrue(grad_op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0", grad_op.get_attr("_XlaScope"))
+
+      # Ensure the ops run: grad(x1*x1) = 2*x1
+      self.assertAllClose([1.0, 1.0, 2.0], sess.run([x, r, g_r]))
+
+  @test_util.build_as_function_and_v1_graph
+  def testPlaysNicelyWithDefunSeparateGradientScope(self):
+    with self.session(graph=ops.Graph()) as sess:
+      with jit.experimental_jit_scope(True):
+
+        @function.Defun(
+            compiled=True, noinline=True, separate_compiled_gradients=True)
+        def mulop(x1, x2):
+          return x1 * x2
+
+        x = constant_op.constant(1.0)
+        r = mulop(x, x)
+        g_r = gradients.gradients(r, x, name="GA")[0]
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s)
+
+      # Ensure the gradient (SymbolicGradient) is compiled, with a different
+      # _XlaScope from the function itself.
+      grad_op = g_r.op.inputs[0].op
+      self.assertTrue(grad_op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0_grad_GA",
+                       grad_op.get_attr("_XlaScope"))
+
+      # Ensure the ops run: grad(x1*x1) = 2*x1
+      self.assertAllClose([1.0, 1.0, 2.0], sess.run([x, r, g_r]))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/pjrt_autoclustering_test.py

@@ -0,0 +1,74 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests single device compilation + autoclustering using the Device API (PjRt).
+
+This feature is still under active development and is protected behind the
+`--tf_xla_use_device_api` flag in the `TF_XLA_FLAGS` environment variable.
+"""
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+
+
+class PjrtAutoclusteringTest(test.TestCase):
+
+  def test_xla_compile_and_run_on_gpu_device(self):
+
+    if not test.is_gpu_available() or not test.is_built_with_gpu_support():
+      test.skipTest("Test only applicable on GPU")
+
+    @def_function.function
+    def arithmetic(x):
+      return 2 * x + 1
+
+    @def_function.function
+    def conditional(x):
+      # cond uses switch and merge, which are not supported by XLA based on
+      # https://docs.google.com/spreadsheets/d/1H8AIDdnlyyaWZOYN3WpBVNOmGOS_M8OyF7IA7kL3fjk/edit?resourcekey=0-I-mIp472YuK8FuBa5Zmzmg#gid=139369773
+      return cond.cond(math_ops.reduce_sum(x) < 5, lambda: x + x, lambda: x)
+
+    @def_function.function
+    def func(x, y):
+      return (arithmetic(x) + conditional(y) ** 2) / 2
+
+    i1 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+
+    # Simple case: all ops supported by XLA
+    with ops.device("/device:GPU:0"):
+      with context.collect_graphs(optimized=True) as graphs:
+        result = arithmetic(i1)
+    self.assertAllClose(result.numpy(), [[3.0, 5.0], [7.0, 9.0]], atol=1e-05)
+    graph_ops = [n.op for n in graphs[0].node]
+    self.assertContainsSubset(["_XlaCompile", "_XlaRun"], graph_ops)
+
+    # Complex case: includes ops not supported by XLA (switch and merge)
+    i2 = constant_op.constant([[5.0, 6.0], [7.0, 8.0]])
+    with ops.device("/device:GPU:0"):
+      with context.collect_graphs(optimized=True) as graphs:
+        result = func(i1, i2)
+    self.assertAllClose(result.numpy(), [[14.0, 20.5], [28, 36.5]], atol=1e-05)
+    graph_ops = [n.op for n in graphs[0].node]
+    self.assertContainsSubset(["_XlaCompile", "_XlaRun"], graph_ops)
+    # because of the cond, not all ops can be combined into a single _XlaCompile
+    self.assertGreater(graph_ops.count("_XlaCompile"), 1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/pjrt_compile_test.py

@@ -0,0 +1,162 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests single device compilation + execution using the Device API (aka PjRt).
+
+This feature is still under active development and is protected behind the
+`--tf_xla_use_device_api` flag in the `TF_XLA_FLAGS` environment variable.
+"""
+import numpy as np
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+
+
+class PjrtCompileTest(test.TestCase):
+
+  def test_compile_on_demand(self):
+    if not test.is_gpu_available() or not test.is_built_with_gpu_support():
+      test.skipTest("Test only applicable on GPU")
+
+    with ops.device("/device:XLA_GPU:0"):
+      a = constant_op.constant([1.0, 2.0])
+      b = constant_op.constant([2.0, 3.0])
+      c = a + b
+      self.assertAllClose([3.0, 5.0], c, atol=1e-05)
+
+      v = variables.Variable([0.0, 1.0])
+      v.assign([1.0, 2.0])
+      self.assertAllClose([1.0, 2.0], v.value(), atol=1e-05)
+      v.assign_add([1.0, 2.0])
+      self.assertAllClose([2.0, 4.0], v.value(), atol=1e-05)
+
+      d = c + v
+      self.assertAllClose([5.0, 9.0], d, atol=1e-05)
+
+  # Tests compilation and execution of a jit_compiled function using PjRt.
+  def test_xla_local_launch(self):
+    if not test.is_gpu_available() or not test.is_built_with_gpu_support():
+      test.skipTest("Test only applicable on GPU")
+
+    @def_function.function(jit_compile=True)
+    def foo(x, y):
+      return x + y + 1
+
+    @def_function.function(jit_compile=True)
+    def bar(x, y):
+      x.assign(y)
+      y.assign_add([1.0, 1.0])
+
+    # Currently PjRt only supports compilation and execution for the XLA_GPU
+    # device to unblock development. Support for non-XLA devices (CPU/GPU/single
+    # core TPU) is going to be added soon, after which support for XLA_* devices
+    # will be dropped.
+    # TODO(b/255826209): Modify the test as we progress towards supporting
+    # non-XLA devices.
+    with ops.device("/device:XLA_GPU:0"):
+      # Function call with scalars
+      self.assertEqual(self.evaluate(foo(1, 2)), 4)
+
+      # Function call with tensors
+      a = constant_op.constant([1.0, 2.0])
+      b = constant_op.constant([2.0, 3.0])
+      self.assertAllClose([4.0, 6.0], foo(a, b), atol=1e-05)
+
+      # Function call with variables
+      x = variables.Variable([0.0, 1.0])
+      y = variables.Variable([1.0, 2.0])
+      self.assertAllClose([2.0, 4.0], foo(x, y), atol=1e-05)
+
+      # Function call with constant and variable
+      self.assertAllClose([2.0, 4.0], foo(a, x), atol=1e-05)
+
+      # Function call that updates variables
+      bar(x, y)
+      self.assertAllClose([1.0, 2.0], x.value(), atol=1e-05)
+      self.assertAllClose([2.0, 3.0], y.value(), atol=1e-05)
+
+  def test_xla_launch_and_tf_kernel_on_gpu_device(self):
+    @def_function.function(jit_compile=True)
+    def const_fn():
+      return constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+
+    @def_function.function(jit_compile=True)
+    def matmul_fn(x):
+      return math_ops.matmul(x, x)
+
+    # The following tests XlaLaunch kernel interoperation with legacy TF
+    # GPU kernel.
+    #
+    # The following use case is tested:
+    # host -> h2d transfer (creating a tf::Tensor with PjRtTensorBuffer)
+    # -> XLA op (GPU) -> TF op (GPU) -> d2h transfer.
+    host_tensor = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+    with ops.device("/device:GPU:0"):
+      xla_tensor = const_fn()
+      xla_result = matmul_fn(host_tensor)
+      result = math_ops.matmul(xla_result, xla_tensor)  # TF GPU kernel.
+
+    # Using numpy to calculate the reference result ("ref_*").
+    ref_tensor = np.array([[1.0, 2.0], [3.0, 4.0]])
+    ref_result = np.matmul(np.matmul(ref_tensor, ref_tensor), ref_tensor)
+    self.assertAllClose(result.numpy(), ref_result, atol=1e-05)
+
+    # The following use case is tested:
+    # host -> h2d transfer (creating a tf::Tensor with PjRtTensorBuffer)
+    # -> TF op (GPU, creating a tf::Tensor with device mem)
+    # -> XLA op (GPU, accepting a tf::Tensor with device mem) -> d2h transfer.
+    with ops.device("/device:GPU:0"):
+      tf_matmul_tensor = math_ops.matmul(host_tensor, host_tensor)
+      xla_result = matmul_fn(tf_matmul_tensor)
+
+    ref_matmul_tensor = np.matmul(ref_tensor, ref_tensor)
+    ref_result_2 = np.matmul(ref_matmul_tensor, ref_matmul_tensor)
+    self.assertAllClose(xla_result.numpy(), ref_result_2, atol=1e-05)
+
+  def test_xla_launch_with_var_on_gpu_device(self):
+    @def_function.function(jit_compile=True)
+    def foo(x, y):
+      return x + y + 1
+
+    @def_function.function(jit_compile=True)
+    def bar(x, y):
+      x.assign(y)
+      y.assign_add([1.0, 1.0])
+
+    with ops.device("/device:GPU:0"):
+      a = constant_op.constant([1.0, 2.0])
+      x = variables.Variable([0.0, 1.0])
+      result_tensor = foo(x, a)
+
+    self.assertAllClose(result_tensor.numpy(), [2.0, 4.0], atol=1e-05)
+
+    # The following use case is tested:
+    # Variable updates following an XLA computation that reads the updated
+    # variables.
+    with ops.device("/device:GPU:0"):
+      var_a = variables.Variable([0.0, 1.0])
+      var_b = variables.Variable([1.0, 2.0])
+      bar(var_a, var_b)
+      result = foo(var_a, var_b)
+
+    self.assertAllClose([1.0, 2.0], var_a.value(), atol=1e-05)
+    self.assertAllClose([2.0, 3.0], var_b.value(), atol=1e-05)
+    self.assertAllClose(result, [4.0, 6.0], atol=1e-05)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/xla_test.py

@@ -0,0 +1,337 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Tests for python.compiler.xla.xla."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.compiler.xla import xla
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+from tensorflow.python.summary import summary
+from tensorflow.python.tpu import tpu_feed
+
+
+_EXPECTED_LOSS = 1
+_EXPECTED_FEATURE = 2
+_EXPECTED_LABEL = 3
+
+
+class XLACompileContextTest(test.TestCase, parameterized.TestCase):
+
+  def create_test_xla_compile_context(self):
+    computation_name = ops.get_default_graph().unique_name('computation')
+    pivot = control_flow_ops.no_op(name=computation_name + '/pivot')
+    return xla.XLACompileContext(name=computation_name, pivot=pivot)
+
+  @test_util.run_v1_only('Testing graph mode behavior only')
+  def test_report_unsupported_operations_graph_mode(self):
+    """Tests that unsupported operations are detected."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    dummy_tensor = constant_op.constant(1.1)
+    audio_summary = summary.audio('audio_summary', dummy_tensor, 0.5)
+    histogram_summary = summary.histogram('histogram_summary', dummy_tensor)
+    image_summary = summary.image('image_summary', dummy_tensor)
+    scalar_summary = summary.scalar('scalar_summary', dummy_tensor)
+    tensor_summary = summary.tensor_summary('tensor_summary', dummy_tensor)
+    summary.merge(
+        [
+            audio_summary, histogram_summary, image_summary, scalar_summary,
+            tensor_summary
+        ],
+        name='merge_summary')
+    logging_ops.Print(dummy_tensor, [dummy_tensor], name='print_op')
+    context.Exit()
+
+    unsupported_ops_names = [op.name for op in context._unsupported_ops]
+    self.assertEqual(unsupported_ops_names, [
+        u'audio_summary', u'histogram_summary', u'image_summary',
+        u'scalar_summary', u'tensor_summary', u'merge_summary/merge_summary',
+        u'print_op'
+    ])
+
+  @test_util.run_v1_only('Testing graph mode behavior only')
+  def test_resource_variable_graph_mode(self):
+    """Tests that resource variable usage is allowed."""
+    a = variable_scope.get_variable(
+        name='variable_a', use_resource=True, initializer=1)
+
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    a.assign(2)
+    context.Exit()
+
+  def test_resource_variable_in_function(self):
+    """Tests that resource variable usage is allowed."""
+    a = variable_scope.get_variable(
+        name='variable_a', use_resource=True, initializer=1)
+
+    @def_function.function
+    def func():
+      context = self.create_test_xla_compile_context()
+      context.Enter()
+      o = a.assign(2)
+      context.Exit()
+      return o
+
+    self.assertEqual(self.evaluate(func()), 2)
+
+  @test_util.run_v1_only('Testing v1-only ref variable handling.')
+  def test_non_resource_variable_error(self):
+    """Tests that non-resource variable usage is disallowed."""
+    a = variable_scope.get_variable(
+        name='variable_a', shape=(1), use_resource=False)
+
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    with self.assertRaisesRegex(
+        NotImplementedError, 'Non-resource Variables are not supported inside '
+        r'XLA computations \(operator name: Assign\)'):
+      state_ops.assign(a, a + 1)
+    context.Exit()
+
+  @test_util.build_as_function_and_v1_graph
+  def test_nested_xla_compile_error(self):
+    """Tests that nested XLA computation leads to fatal error."""
+    context1 = self.create_test_xla_compile_context()
+    context1.Enter()
+
+    context2 = self.create_test_xla_compile_context()
+    context2.Enter()
+    with self.assertRaisesRegex(ValueError,
+                                'XLA compiled computations cannot be nested'):
+      constant_op.constant(1)
+    context2.Exit()
+    context1.Exit()
+
+  @test_util.build_as_function_and_v1_graph
+  def test_xla_compile_attr(self):
+    """Tests that ops are tagged with XLA compile ID attribute."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+    self.assertIn('_xla_compile_id', op.op.node_def.attr)
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_without_input(self):
+    """Tests that ops without inputs depend on pivot correctly."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+
+    self.assertIn(context._pivot, op.op.control_inputs)
+
+  @test_util.run_v1_only('Testing graph mode behavior only')
+  def test_external_control_edges_graph_mode(self):
+    """Tests that external control edges are handled correctly."""
+    i = constant_op.constant(1)
+    op1 = constant_op.constant(1)
+
+    with ops.control_dependencies([op1]):
+      op2 = constant_op.constant(1)
+    self.assertIn(op1.op, op2.op.control_inputs)
+
+    def while_body(i):
+      del i  # unused
+      context = self.create_test_xla_compile_context()
+      context.Enter()
+      with ops.control_dependencies([op1]):
+        op3 = constant_op.constant(1)
+      context.Exit()
+      self.assertNotIn(op1.op, op3.op.control_inputs)
+      return op3
+
+    while_loop.while_loop(
+        cond=lambda i: math_ops.less(i, 10), body=while_body, loop_vars=[i])
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_output_marked_as_seen(self):
+    """Tests that any op output is marked as seen in context."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+
+    self.assertIn(op.name, context._values)
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_is_in_context(self):
+    """Tests that XLACompileContext is recognized as an XLA context."""
+    op1 = constant_op.constant(1)
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op2 = constant_op.constant(2)
+    context.Exit()
+    self.assertFalse(control_flow_util.IsInXLAContext(op1.op))
+    self.assertTrue(control_flow_util.IsInXLAContext(op2.op))
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_prevent_feeding(self):
+    """Tests that ops created inside XLACompileContext can not be fed."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+    self.assertFalse(op.graph.is_feedable(op.op))
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_prevent_fetching(self):
+    """Tests that ops created inside XLACompileContext can not be fetched."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+    self.assertFalse(op.graph.is_fetchable(op.op))
+
+
+class XlaCompileTest(test.TestCase):
+
+  @test_util.run_v2_only
+  @test_util.disable_tfrt(
+      'Legacy XLA test. It depends on EncapsulateXlaComputationsPass.')
+  def test_xla_compile_eager(self):
+    """Tests that xla.compile raises proper exception when used eagerly."""
+
+    def computation(a, b):
+      return a + b
+
+    self.assertEqual(self.evaluate(xla.compile(computation, [1, 2])[0]), 3)
+
+  @test_util.disable_tfrt(
+      'Legacy XLA test. It depends on EncapsulateXlaComputationsPass.')
+  def test_xla_compile_in_function(self):
+    """Tests that xla.compile works in tf.function."""
+
+    @def_function.function
+    def func_wrapper(a):
+
+      def compute(a):
+        return a + 1
+
+      return xla.compile(compute, [a])
+
+    self.assertEqual(self.evaluate(func_wrapper(1))[0], 2)
+
+  @test_util.disable_tfrt(
+      'Legacy XLA test. It depends on EncapsulateXlaComputationsPass.')
+  def test_xla_compile_write_variable_in_function(self):
+    """Tests that xla.compile works with variable in tf.function."""
+    a = variable_scope.get_variable(
+        name='variable_a', use_resource=True, initializer=1)
+
+    @def_function.function
+    def func_wrapper():
+
+      def compute():
+        a.assign_add(1)
+        a.assign_sub(2)
+        return a.read_value()
+
+      return xla.compile(compute)
+
+    self.evaluate(a.initializer)
+    self.assertEqual(self.evaluate(func_wrapper())[0], 0)
+
+
+class CheckFunctionArgumentCountTest(test.TestCase):
+
+  def test_simple(self):
+    """Tests that arg checker works for functions with no varargs or defaults.
+    """
+
+    def func(x, y, z):
+      return x + y + z
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual('exactly 3 arguments',
+                     xla.check_function_argument_count(func, 2, None))
+    queue = tpu_feed.InfeedQueue(2)
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual('exactly 3 arguments',
+                     xla.check_function_argument_count(func, 2, queue))
+
+  def test_default_args(self):
+    """Tests that arg checker works for a function with no varargs."""
+
+    def func(x, y, z=17):
+      return x + y + z
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, None))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 1, None))
+    self.assertEqual('at most 3 arguments',
+                     xla.check_function_argument_count(func, 4, None))
+    queue = tpu_feed.InfeedQueue(1)
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 0, queue))
+    self.assertEqual('at most 3 arguments',
+                     xla.check_function_argument_count(func, 4, queue))
+
+  def test_var_args(self):
+    """Tests that arg checker works for a function with varargs."""
+
+    def func(x, y, *z):
+      return x + y + len(z)
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 4, None))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 1, None))
+    queue = tpu_feed.InfeedQueue(1)
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, queue))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 0, queue))
+
+  def test_var_args_and_defaults(self):
+    """Tests that arg checker works for a function with varargs and defaults."""
+
+    def func(x, y, z=17, *q):  # pylint: disable=keyword-arg-before-vararg
+      return x + y + z + len(q)
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 4, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 5, None))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 1, None))
+    queue = tpu_feed.InfeedQueue(1)
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 4, queue))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 0, queue))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/compiler/xla/xla.py

@@ -0,0 +1,620 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""xla is an experimental library that provides XLA support APIs."""
+
+import contextlib
+
+
+from tensorflow.compiler.jit.ops import xla_ops
+from tensorflow.compiler.jit.ops import xla_ops_grad  # pylint: disable=unused-import
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.distribute import summary_op_util
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util.compat import collections_abc
+from tensorflow.python.util.deprecation import deprecated
+from tensorflow.python.util.tf_export import tf_export
+
+_XLA_COMPILE_ATTR = '_xla_compile_id'
+_MAX_WARNING_LINES = 5
+
+# Operations that indicate some error in the users graph. For example, XLA
+# computation should not have any Placeholder op.
+_DENYLISTED_OPS = set([
+    'Placeholder',
+])
+
+# XLA doesn't currently support reading of intermediate tensors, thus some ops
+# are not supported.
+_UNSUPPORTED_OPS = set([
+    'AudioSummary',
+    'AudioSummaryV2',
+    'HistogramSummary',
+    'ImageSummary',
+    'MergeSummary',
+    'Print',
+    'ScalarSummary',
+    'TensorSummary',
+    'TensorSummaryV2',
+])
+
+
+@tf_export('xla.experimental.compile')
+@deprecated(
+    None, 'xla.experimental.compile is deprecated. Consider using '
+    '`@tf.function(jit_compile=True)`.',
+    warn_once=True)
+def compile(computation, inputs=None):  # pylint: disable=redefined-builtin
+  """Builds an operator that compiles and runs `computation` with XLA.
+
+  NOTE: In eager mode, `computation` will have `@tf.function` semantics.
+
+  Args:
+    computation: A Python function that builds a computation to apply to the
+      input. If the function takes n inputs, 'inputs' should be a list of n
+      `Tensor`s.
+
+      `computation` may return a list of `Tensor`s and `Operation`s.
+      `Tensor`s must come before `Operation`s in the returned list.
+
+      All `Operation`s returned from `computation` will be executed when
+      evaluating any of the returned output tensors.
+    inputs: A list of inputs or `None` (equivalent to an empty list). Each input
+      can be a nested structure containing values that can be converted to
+      `Tensor`s. Note that passing an N-dimension list of compatible values will
+      result in an N-dimension list of scalar `Tensor`s rather than a single
+      Rank-N `Tensor`. If you need a different behavior, convert parts of
+      `inputs` to `Tensor`s with `tf.convert_to_tensor`.
+
+  Returns:
+    List of `Tensor`s corresponding to the `Tensor`s from
+      the output of `computation` i.e. the same return value as if
+      computation(*inputs) is called directly, with the following exceptions:
+      * None output: a NoOp would be returned with a control dependency on
+         `computation`.
+      * Single value output: a tuple containing the value would be returned.
+      * Operation-only outputs: a NoOp would be returned with a control
+      dependency on `computation`.
+      TODO(b/121383831): Investigate into removing these special cases.
+
+  Raises:
+    RuntimeError: When eager execution is enabled.
+
+  Known issues:
+    When a tf.random operation is built with XLA, the implementation doesn't
+      pass the user provided seed to the XLA compiler. As such, the XLA compiler
+      generates a random number and uses it as a seed when compiling the
+      operation. This implementation causes a violation of the Tensorflow
+      defined semantics in two aspects. First, changing the value of the user
+      defined seed doesn't change the numbers generated by the operation.
+      Second, when a seed is not specified, running the program multiple times
+      will generate the same numbers.
+  """
+  if context.executing_eagerly():
+
+    @def_function.function
+    def xla_compile_wrapper():
+      return _compile_internal(computation, inputs)
+
+    return xla_compile_wrapper()
+
+  return _compile_internal(computation, inputs)
+
+
+class XLACompileContext(control_flow_ops.XLAControlFlowContext):
+  """A `ControlFlowContext` for nodes inside an XLA computation cluster.
+
+  THIS IS ONLY FOR TENSORFLOW INTERNAL IMPLEMENTATION, DO NO USE DIRECTLY.
+
+  The primary role of `XLACompileContext` is to mark operators inside a
+  xla.compile() computation with attribute "_xla_compile_id=XYZ", where XYZ is
+  a unique name.
+
+  `ControlFlowContext` is used to perform the annotation since it integrates
+  with Tensorflow constructs like ResourceVariables. For example, if a
+  `ResourceVariable` is constructed inside a xla.compile() block, the
+  `ResourceVariable` implementation can use
+  `with ops.control_dependencies(None)` to build the variable's definition
+  outside the compiled computation.
+  """
+
+  def __init__(self, name, pivot):
+    """Builds a new XLACompileContext.
+
+    Args:
+      name: a unique name for the context, used to populate the
+        `_xla_compile_id` attribute.
+      pivot: a pivot node. Nodes in the XLACompileContext that do not have any
+        inputs will have a control dependency on the pivot node. This ensures
+        that nodes are correctly included in any enclosing control flow
+        contexts.
+    """
+    super(XLACompileContext, self).__init__()
+    self._name = name
+    self._name_as_bytes = compat.as_bytes(name)
+    self._unsupported_ops = []
+    self._pivot = pivot
+
+  def report_unsupported_operations(self):
+    if self._unsupported_ops:
+      op_str = '\n'.join([
+          '  %s (%s)' % (op.type, op.name)
+          for op in self._unsupported_ops[:_MAX_WARNING_LINES]
+      ])
+      logging.warning('%d unsupported operations found: \n%s',
+                      len(self._unsupported_ops), op_str)
+      if len(self._unsupported_ops) > _MAX_WARNING_LINES:
+        logging.warning('... and %d more',
+                        len(self._unsupported_ops) - _MAX_WARNING_LINES)
+
+  def _RemoveExternalControlEdges(self, op: ops.Operation):
+    """Remove any external control dependency on this op."""
+    internal_control_inputs = []
+    external_control_inputs = []
+    for x in op.control_inputs:
+      # pylint: disable=protected-access
+      is_internal_op = False
+      ctxt = x._get_control_flow_context()
+      while ctxt is not None:
+        if ctxt == self:
+          is_internal_op = True
+          break
+        ctxt = ctxt._outer_context
+      if is_internal_op:
+        internal_control_inputs.append(x)
+      else:
+        external_control_inputs.append(x)
+      # pylint: enable=protected-access
+    # pylint: disable=protected-access
+    op._remove_all_control_inputs()
+    op._add_control_inputs(internal_control_inputs)
+    # pylint: enable=protected-access
+    return internal_control_inputs, external_control_inputs
+
+  def AddOp(self, op: ops.Operation):
+    """Create op in XLACompileContext and notifies outer context recursively."""
+    # pylint: disable=protected-access
+    if op.type in _DENYLISTED_OPS:
+      logging.error(
+          'Operation of type %s (%s) is not supported in XLA. Execution will '
+          'fail if this op is used in the graph. ', op.type, op.name)
+
+    # TODO(ycao): Automatically disable summaries instead of reporting them.
+    if op.type in _UNSUPPORTED_OPS:
+      self._unsupported_ops.append(op)
+
+    if any(x.dtype._is_ref_dtype for x in op.inputs):
+      raise NotImplementedError(
+          'Non-resource Variables are not supported inside XLA computations '
+          '(operator name: %s)' % op.name)
+
+    if _XLA_COMPILE_ATTR in op.node_def.attr:
+      raise ValueError('XLA compiled computations cannot be nested, (operator '
+                       'name: %s)' % op.name)
+
+    op._set_attr(
+        _XLA_COMPILE_ATTR, attr_value_pb2.AttrValue(s=self._name_as_bytes))
+
+    op.graph.prevent_feeding(op)
+    op.graph.prevent_fetching(op)
+
+    # Remove any control edges from outer control flow contexts. These may cause
+    # mismatched frame errors. An example is when one of op's inputs is
+    # generated in a different While control flow context.
+    (internal_control_inputs,
+     external_control_inputs) = self._RemoveExternalControlEdges(op)
+
+    if not op.inputs:
+      # Add a control edge from the control pivot to this op.
+      if not internal_control_inputs:
+        # pylint: disable=protected-access
+        op._add_control_input(self._pivot)
+        # pylint: enable=protected-access
+    else:
+      for index in range(len(op.inputs)):
+        x = op.inputs[index]
+        real_x = self.AddValue(x)
+        if real_x is not x:
+          op._update_input(index, real_x)  # pylint: disable=protected-access
+
+    if external_control_inputs:
+      # Use an identity to pull control inputs as data inputs. Note that we
+      # ignore ops which don't have outputs. TODO(phawkins): fix that.
+      with ops.control_dependencies(None):
+        self.Enter()
+        external_control_inputs = [
+            array_ops.identity(x.outputs[0]).op
+            for x in external_control_inputs
+            if x.outputs
+        ]
+        self.Exit()
+      # pylint: disable=protected-access
+      op._add_control_inputs(external_control_inputs)
+      # pylint: enable=protected-access
+
+    # Mark op's outputs as seen by this context and any outer contexts.
+    output_names = [x.name for x in op.outputs]
+    context = self
+    while context is not None:
+      # pylint: disable=protected-access
+      context._values.update(output_names)
+      context = context._outer_context
+      # pylint: enable=protected-access
+
+    if self._outer_context:
+      self._outer_context.AddInnerOp(op)
+
+  def AddValue(self, val):
+    """Add `val` to the current context and its outer context recursively."""
+    if val.name in self._values:
+      # Use the real value if it comes from outer context.
+      result = self._external_values.get(val.name)
+      return val if result is None else result
+
+    result = val
+    self._values.add(val.name)
+    if self._outer_context:
+      result = self._outer_context.AddValue(val)
+      self._values.add(result.name)
+
+    self._external_values[val.name] = result
+
+    return result
+
+  def AddInnerOp(self, op: ops.Operation):
+    self.AddOp(op)
+    if self._outer_context:
+      self._outer_context.AddInnerOp(op)
+
+  @property
+  def grad_state(self):
+    # Define the gradient loop state associated with the XLACompileContext to
+    # be None as the XLACompileContext does not get nested nor does the
+    # grad_state outside the XLACompileContext affect the graph inside so the
+    # grad_state should be as if this is the top-level gradient state.
+    return None
+
+  @property
+  def back_prop(self):
+    """Forwards to the enclosing while context, if any."""
+    if self.GetWhileContext():
+      return self.GetWhileContext().back_prop
+    return False
+
+
+def _compile_internal(computation, inputs=None):
+  """Builds graph operators that compiles and symbolically executes computation.
+
+  Args:
+    computation: A Python function that builds the computation to compile and
+      execute.
+    inputs: A list of inputs or `None` (equivalent to an empty list). Each input
+      can be a nested structure containing values that are convertible to
+      tensors. Note that passing an N-dimension list of compatible values will
+      result in a N-dimension list of scalar tensors rather than a single Rank-N
+      tensors. If you need different behavior, convert part of inputs to tensors
+      with `tf.convert_to_tensor`.
+
+  Returns:
+    Same data structure as if computation(*inputs) is called directly with some
+    exceptions for correctness. Exceptions include: 1) None output 2) Single
+    value output 3) Operation-only outputs
+  Raises:
+    ValueError: If any element in computation outputs is neither an operations
+      or a value that can be converted to tensor.
+    ValueError: If computation outputs is non-flat and contains any Operations.
+    TypeError: If `inputs` is not a list or tuple.
+  """
+  if inputs is None:
+    inputs = []
+
+  if not isinstance(inputs, collections_abc.Sequence):
+    raise TypeError('inputs must be a list')
+
+  # Flatten inputs.
+  flat_inputs = nest.flatten(inputs)
+  # Converts inputs to Tensors.
+  flat_inputs = [ops.convert_to_tensor(x) for x in flat_inputs]
+
+  cluster_name = ops.get_default_graph().unique_name('cluster')
+  pivot = control_flow_ops.no_op(name=cluster_name + '/pivot')
+  context = XLACompileContext(name=cluster_name, pivot=pivot)
+  try:
+    context.Enter()
+
+    # Add identity ops so even unused inputs are 'consumed' by the
+    # computation.
+    flat_inputs = [
+        array_ops.identity(x, name='input_{}'.format(i))
+        for i, x in enumerate(flat_inputs)
+    ]
+
+    # Re-pack flat_inputs in same structure as 'inputs'.
+    computation_inputs = nest.pack_sequence_as(
+        structure=inputs, flat_sequence=flat_inputs)
+
+    # Only resource variables work inside an XLA computation, so turn on
+    # resource variables for the computation.
+    vscope = variable_scope.get_variable_scope()
+    saved_use_resource = vscope.use_resource
+    vscope.set_use_resource(True)
+
+    with _disable_summary_context():
+      outputs = computation(*computation_inputs)
+
+    # Restore variable scope after computation.
+    vscope.set_use_resource(saved_use_resource)
+
+    outputs_is_flat = is_flat(outputs)
+    if outputs_is_flat:
+      output_tensors, control_deps = _postprocess_flat_outputs(outputs)
+    else:
+      output_tensors, control_deps = _postprocess_non_flat_outputs(outputs)
+
+    context.ExitResult(output_tensors)
+  finally:
+    context.report_unsupported_operations()
+    context.Exit()
+
+  # When XLA computation returns only operations and no tensors, a NoOp
+  # dependent on the operations in outputs is returned. Otherwise final
+  # outputs would be empty and there is no way to trigger returned
+  # operations.
+  if not output_tensors:
+    return control_flow_ops.group(control_deps, name='output_0')
+
+  output_tensors = [
+      xla_ops.xla_cluster_output(o, name='output{}'.format(i))
+      for i, o in enumerate(output_tensors)
+  ]
+
+  with ops.control_dependencies(control_deps):
+    # Wraps the outputs in identity operators that carries control
+    # dependencies.
+    output_tensors = [
+        array_ops.identity(o, name='output_%d' % i)
+        for i, o in enumerate(output_tensors)
+    ]
+
+  # If `computation` returned non-flat output structure, pack output tensors
+  # back into same structure.
+  if not outputs_is_flat:
+    output_tensors = nest.pack_sequence_as(
+        structure=outputs, flat_sequence=output_tensors)
+
+  return output_tensors
+
+
+def is_flat(outputs):
+  """Checks if outputs is a flat structure.
+
+    Following structures and values are considered flat:
+    1) None
+    2) A single object
+    3) A list or tuple of Tensors/Operations
+
+    The only structures that this function understands are sequences,
+    dictionaries and types defined using the attrs library.  E.g. this means
+    that if outputs contains a single user-defined Object, it is considered to
+    be flat. Errors are raised later on if that Object cannot be converted to a
+    Tensor.
+
+  Args:
+    outputs: Output from `computation` inside `xla.compile`.
+
+  Returns:
+    A boolean indicates whether outputs is flat.
+  """
+  # If outputs is a list or tuple, check if it has any nested structure. If
+  # there is, then outputs is non-flat.
+  if isinstance(outputs, collections_abc.Sequence):
+    for o in outputs:
+      if (isinstance(o, collections_abc.Sequence) or
+          isinstance(o, collections_abc.Mapping) or
+          hasattr(o.__class__, '__attrs_attrs__')):
+        return False
+
+  # If outputs is a dict, it is non-flat.
+  if isinstance(outputs, collections_abc.Mapping):
+    return False
+
+  # If outputs is from the attrs library, it is non-flat.
+  if hasattr(outputs.__class__, '__attrs_attrs__'):
+    return False
+
+  # Getting here means either outputs itself is a single non-structured value
+  # or it is a flat list of single non-structured values.
+  return True
+
+
+def _postprocess_flat_outputs(outputs):
+  """Validates flat outputs and adds back device assignments.
+
+  Args:
+    outputs: Output from `computation` inside `xla.compile`.
+
+  Returns:
+    Tensors and Operations extracted from outputs.
+  """
+  # Following code segment is to preserve legacy behavior. Previously we only
+  # supported flat outputs and thus for consistency it was nice to convert even
+  # single element into a tuple. But now that we support arbitrary output
+  # structure, this is no longer necessary.
+  # TODO(b/121383831): Migrate all legacy use cases and delete this special
+  # case.
+  # If the computation returns `None`, make it an empty tuple.
+  if outputs is None:
+    outputs = tuple()
+  # If the computation only returned one value, make it a tuple.
+  if not isinstance(outputs, collections_abc.Sequence):
+    outputs = (outputs,)
+
+  # Append `no_op` here so that return value of this function always contains
+  # at least one op that can trigger XlaLaunch node.
+  outputs += (control_flow_ops.no_op(),)
+  try:
+    outputs = [
+        o if isinstance(o, ops.Operation) else ops.convert_to_tensor(o)
+        for o in outputs
+    ]
+  except Exception as e:
+    raise ValueError(
+        'XLA computation function return values must all either be Operations'
+        ' or convertible to Tensors. Got error: "%s"' % str(e))
+
+  # Separates the returned Operations and Tensors.
+  output_operations = [o for o in outputs if isinstance(o, ops.Operation)]
+  output_tensors = [o for o in outputs if not isinstance(o, ops.Operation)]
+
+  if outputs != output_tensors + output_operations:
+    raise ValueError(
+        'XLA computation function must return zero or more Tensor values '
+        'followed by zero or more Operations.')
+
+  new_output_tensors = []
+  for t in output_tensors:
+    with ops.device(t.device if t.device else ''):
+      new_output_tensors.append(array_ops.identity(t))
+
+  return new_output_tensors, output_operations
+
+
+def _postprocess_non_flat_outputs(outputs):
+  """Validates non-flat outputs and adds back device assignments.
+
+  Args:
+    outputs: Output from `computation` inside `xla.compile`.
+
+  Returns:
+    Tensors extracted from outputs and an empty list because Operations are not
+    allowed in non-flat outputs..
+  """
+  # Convert all non-Operation outputs to Tensors.
+  new_output_tensors = []
+  for o in nest.flatten(outputs):
+    if isinstance(o, ops.Operation):
+      raise ValueError(
+          'xla.compile does not support Operation as return value in non-flat '
+          'output structure. You can set returned Operations as control '
+          'dependencies of returned Tensors so Operations are triggered when '
+          'Tensors are evaluated. Operation found: "%s"' % o.name)
+
+    try:
+      o = ops.convert_to_tensor(o)
+    except Exception as e:
+      raise ValueError(
+          'XLA computation function return values must all either be '
+          'Operations or convertible to Tensors. Got error: "%s"' % str(e))
+
+    # Makes sure even pass-through inputs/outputs are touched in compile
+    # context by creating an Identity node inside compile context.
+    with ops.device(o.device if o.device else ''):
+      new_output_tensors.append(array_ops.identity(o))
+
+  return new_output_tensors, []
+
+
+@contextlib.contextmanager
+def _disable_summary_context():
+  """Enters a context where all summary ops are skipped.
+
+  Summaries are not yet supported in xla.compile(). So we provide this context
+  manager that can skip creating summary ops. This is a temporary workaround due
+  to XLA not supporting summary ops.
+
+  Yields:
+    None.
+  """
+  original_skip_summary_func = summary_op_util.skip_summary
+  summary_op_util.skip_summary = lambda: True
+
+  try:
+    yield
+  finally:
+    summary_op_util.skip_summary = original_skip_summary_func
+
+
+class _CapturedObject(object):
+  """A placeholder to capture an object."""
+
+  def __init__(self):
+    self._object = None
+
+  def capture(self, o):
+    if self._object:
+      raise RuntimeError(
+          'InternalError: _CapturedObject can capture only once. Please file '
+          'bug.')
+
+    self._object = o
+
+  def get(self):
+    return self._object
+
+
+def check_function_argument_count(func, input_arity, infeed_queue):
+  """Validate the number of input arguments to an XLA function.
+
+  Args:
+    func: the Python function that will be called to generate the body of an XLA
+      computation graph.
+    input_arity: the number of explicit arguments supplied by the caller.
+    infeed_queue: if not None, the infeed queue that will supply
+      additional arguments to the function.
+
+  Returns:
+    None if function can be called with the supplied number of
+      arguments, or an error string if it cannot.
+  """
+  def format_error(complaint, quantity):
+    return '%s %d argument%s' % (complaint, quantity, ''
+                                 if quantity == 1 else 's')
+
+  num_args_supplied = input_arity
+  if infeed_queue is not None:
+    num_args_supplied += infeed_queue.number_of_tuple_elements
+  arg_spec = tf_inspect.getargspec(func)
+  num_func_args = len(arg_spec.args)
+  if arg_spec.defaults is None:
+    num_func_defaults = 0
+  else:
+    num_func_defaults = len(arg_spec.defaults)
+  min_func_args = num_func_args - num_func_defaults
+  if num_args_supplied < min_func_args:
+    # The required number of arguments is not enough to call the function.
+    if num_func_defaults == 0 and arg_spec.varargs is None:
+      return format_error('exactly', num_func_args)
+    else:
+      return format_error('at least', min_func_args)
+  if arg_spec.varargs is None and num_args_supplied > num_func_args:
+    # The required number of arguments is too many to call the function.
+    if num_func_defaults == 0:
+      return format_error('exactly', num_func_args)
+    else:
+      return format_error('at most', num_func_args)
+  # Reaching here means either
+  # 1) There are varargs, func can accept any number of arguments greater than
+  # the minimum.
+  # 2) Number of supplied arguments falls in range of acceptable argument count
+  # of func.
+  return None