chore: import upstream snapshot with attribution

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()