chore: import upstream snapshot with attribution

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