chore: import upstream snapshot with attribution

tensorflow/python/distribute/experimental/BUILD

@@ -0,0 +1,199 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load(
+    "//tensorflow/dtensor:build_defs.bzl",
+    "GPU_2DEVS_BACKEND",
+    "TPU_V3_DONUT_BACKEND",
+    "dtensor_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//learning/brain/distribute/python/benchmark:__subpackages__",
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "experimental",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":mirrored_strategy",
+        ":multi_worker_mirrored_strategy",
+    ],
+)
+
+py_library(
+    name = "mirrored_strategy",
+    srcs = ["mirrored_strategy.py"],
+    strict_deps = True,
+    deps = [
+        ":dtensor_strategy_extended",
+        ":dtensor_util",
+        "//tensorflow/dtensor/python:config",
+        "//tensorflow/dtensor/python:mesh_util",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/framework:device",
+    ],
+)
+
+dtensor_test(
+    name = "mirrored_strategy_test",
+    srcs = ["mirrored_strategy_test.py"],
+    shard_count = {"tpu": 2},
+    tags = ["no_pip"],
+    deps = [
+        ":dtensor_util",
+        ":mirrored_strategy",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:d_variable",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/dtensor/python:mesh_util",
+        "//tensorflow/dtensor/python/tests:test_util",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:stateless_random_ops",
+        "//tensorflow/python/ops:variables",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "dtensor_util",
+    srcs = ["dtensor_util.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/dtensor/python:accelerator_util",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:input_util",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/python/distribute:cross_device_ops",
+        "//tensorflow/python/distribute:device_util",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion_registry",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:summary_ops_v2",
+    ],
+)
+
+dtensor_test(
+    name = "dtensor_util_test",
+    srcs = ["dtensor_util_test.py"],
+    tags = ["no_pip"],
+    deps = [
+        ":dtensor_util",
+        ":mirrored_strategy",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/dtensor/python/tests:test_util",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "dtensor_strategy_extended",
+    srcs = ["dtensor_strategy_extended.py"],
+    strict_deps = True,
+    deps = [
+        ":dtensor_util",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:config",
+        "//tensorflow/dtensor/python:d_variable",
+        "//tensorflow/dtensor/python:input_util",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:distribute_utils",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_library(
+    name = "multi_worker_mirrored_strategy",
+    srcs = ["multi_worker_mirrored_strategy.py"],
+    strict_deps = True,
+    deps = [
+        ":dtensor_strategy_extended",
+        ":dtensor_util",
+        "//tensorflow/dtensor/python:config",
+        "//tensorflow/dtensor/python:mesh_util",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:multi_worker_util",
+        "//tensorflow/python/distribute/cluster_resolver:tfconfig_cluster_resolver_py",
+    ],
+)
+
+dtensor_test(
+    name = "multi_worker_mirrored_strategy_test",
+    srcs = ["multi_worker_mirrored_strategy_test.py"],
+    additional_backends = [
+        GPU_2DEVS_BACKEND,
+        TPU_V3_DONUT_BACKEND,
+    ],
+    disable = [
+        "gpu",  # multi-client gpu is tested via GPU_2DEVS_BACKEND.
+        "tpu",  # multi-client tpu is tested via TPU_V3_DONUT_BACKEND.
+    ],
+    disable_tfrt = [
+        "cpu",  # TODO(b/217969210): Re-enable in TFRT CPU.
+        GPU_2DEVS_BACKEND,  # TODO(b/230679405): Re-enable in TFRT GPU.
+    ],
+    tags = [
+        "no_pip",
+        "no_windows",
+        "nosan",
+    ],  # b/195537906
+    deps = [
+        ":dtensor_util",
+        ":multi_worker_mirrored_strategy",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:d_variable",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/dtensor/python/tests:multi_client_test_util",
+        "//tensorflow/dtensor/python/tests:test_util",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute/cluster_resolver:tfconfig_cluster_resolver_py",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:stateless_random_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/flags",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/distribute/experimental/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+# ==============================================================================
+"""Experimental Distribution Strategy library."""
+
+# pylint: disable=unused-import
+from tensorflow.python.distribute.experimental import mirrored_strategy
+from tensorflow.python.distribute.experimental import multi_worker_mirrored_strategy
+# pylint: enable=unused-import

tensorflow/python/distribute/experimental/dtensor_strategy_extended.py

@@ -0,0 +1,278 @@
+# 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.
+# ==============================================================================
+"""Implement a StrategyExtended based on the DTensor low level API."""
+
+import functools
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import config as d_config
+from tensorflow.dtensor.python import d_variable
+from tensorflow.dtensor.python import input_util
+from tensorflow.dtensor.python import layout
+from tensorflow.python.data.experimental.ops import distribute
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.util import nest
+
+
+class DTensorStrategyExtended(distribute_lib.StrategyExtendedV2):
+  """Strategy extension that support both single and multi worker strategy."""
+  # Note that the unit test for this class is via the strategy interface.
+
+  def __init__(self, container_strategy, mesh):
+    super().__init__(container_strategy)
+    self._mesh = mesh
+    self._num_clients = d_config.num_clients()
+    self._client_id = d_config.client_id()
+
+  def _create_variable(self, next_creator, **kwargs):
+    # Make sure the pop the `use_resource` which is not supported by the
+    # base tf.Variable. The `use_resource` is added by
+    # creator_with_resource_vars in distribute_lib.py
+    kwargs.pop('use_resource', None)
+
+    # Ignore the colocate_with for the mirrored strategy. Each of the device
+    # will get same copy of variable in the DTensor's case.
+    # `colocate_with` is added when user call:
+    # strategy.extended.colocate_vars_with(variable)
+    kwargs.pop('colocate_with', None)
+
+    # Ignore expected_shape, which is from the v1 Variable. Keras was somehow
+    # using the v1 Variable, but didn't specify that value particularly.
+    kwargs.pop('expected_shape', None)
+
+    # Make sure to call DVariable initializer under the scope so that it will
+    # have the proper replicated layout. The initial_value is multi-typed,
+    # eg it can be a tensor, or a python/numpy type, or a callable that
+    # produce tensor/python/numpy types. In all those cases, we need to wrap
+    # them invoke convert_to_tensor() under the scope so that the proper
+    # layout can be assigned.
+
+    # TODO(scottzhu): The layout information should be injected via kwargs, or
+    # lazily set later.
+    initial_value = kwargs.pop('initial_value')
+    dtype = kwargs.get('dtype', None)
+    def new_initial_value():
+      if callable(initial_value):
+        init_var = ops.convert_to_tensor(initial_value(), dtype=dtype)
+      else:
+        init_var = ops.convert_to_tensor(initial_value, dtype=dtype)
+      rank = init_var.shape.rank
+      return d_api.copy_to_mesh(
+          init_var, layout.Layout.replicated(self._mesh, rank))
+
+    return d_variable.DVariable(new_initial_value, **kwargs)
+
+  @property
+  def _num_replicas_in_sync(self):
+    # The mesh should be 1D with batch sharding only.
+    # In the model parallel case, it should only return the size of
+    # batch dimension.
+    return self._mesh.size
+
+  def value_container(self, value):
+    return value
+
+  @property
+  def worker_devices(self):
+    # Note that in either single worker (MirroredStrategy) or multi worker (
+    # MultiWorkerMirroredStrategy), worker_devices refers to the local worker
+    # devices.
+    return tuple(self._mesh.local_devices())
+
+  @property
+  def parameter_devices(self):
+    # Same as the worker_devices.
+    return self.worker_devices
+
+  def _in_multi_worker_mode(self):
+    return d_config.num_clients() > 1
+
+  def _get_local_replica_id(self, replica_id_in_sync_group):
+    return replica_id_in_sync_group
+
+  def _default_device_scope(self):
+    return d_api.default_mesh(self._mesh)
+
+  def _experimental_distribute_dataset(self, dataset, options):
+    # Strategy always assume the user input data is a batched dataset for
+    # experimental_distribute_dataset().
+    # TODO(yuefengz): Add check for whether a dataset is batched for all
+    # strategies.
+
+    # TODO(b/265198795): Support dataset already batched to global batch size.
+    # Since DTensorDataset doesn't support batched dataset that is already
+    # batched global batch size, it only supports dataset that is batched to
+    # local batch size, we need to infer the batch size, and unbatch the dataset
+    # until the b/265198795 is resolved.
+    batch_size = distribute.compute_batch_size(dataset)
+
+    # There are multiple case that the batch is not static, eg partial batch,
+    # or uneven batch, in all those case, it will return -1.
+    if batch_size.numpy() < 0:
+      # When we don't have a static batch size.
+      raise ValueError('DTensor strategy requires a static batch size for now.'
+                       'The dynamic batch size will be supported in future')
+    # Unbatch the dataset for now since the DTensorDataset has some limitation
+    # about the local batch size as well as the mesh size.
+    dataset = dataset.unbatch()
+
+    def _create_batch_layout(tensor_spec):
+      # For unbatched dataset, the new layout need to have +1 rank for
+      # the batched result.
+      rank = len(tensor_spec.shape) + 1
+      return layout.Layout.batch_sharded(
+          self._mesh, batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+          rank=rank)
+
+    layouts = nest.map_structure(_create_batch_layout, dataset.element_spec)
+
+    return input_util.DTensorDataset(
+        dataset=dataset,
+        mesh=self._mesh,
+        layouts=layouts,
+        global_batch_size=batch_size,
+        dataset_already_batched=False,
+        batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+        # TODO(scottzhu): Add prefetch support by inspecting the input dataset.
+        prefetch=None,
+        tf_data_service_config=None
+    )
+
+  def _make_dataset_iterator(self, dataset):
+    raise NotImplementedError(
+        'Strategy.make_dataset_iterator() is deprecated, and only available '
+        'in the V1 API.')
+
+  def _make_input_fn_iterator(self, input_fn, replication_mode):
+    raise NotImplementedError(
+        'Strategy.make_input_fn_iterator() is deprecated, and only available '
+        'in the V1 API.')
+
+  def _distribute_datasets_from_function(self, dataset_fn, options):
+    # TODO(scottzhu): Implement the logic for options in future
+    del options
+    input_context = distribute_lib.InputContext(
+        num_input_pipelines=self._num_clients,
+        input_pipeline_id=self._client_id,
+        num_replicas_in_sync=self._num_replicas_in_sync
+    )
+    dataset = dataset_fn(input_context)
+
+    # Note that the dataset should already batched to local per-relica batch
+    def _create_batch_layout(tensor_spec):
+      rank = len(tensor_spec.shape)
+      return layout.Layout.batch_sharded(
+          self._mesh, batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+          rank=rank)
+
+    layouts = nest.map_structure(_create_batch_layout, dataset.element_spec)
+
+    batch_size = distribute.compute_batch_size(dataset)
+    # There are multiple case that the batch is not static, eg partial batch,
+    # or uneven batch, in all those case, it will return -1.
+    if batch_size.numpy() < 0:
+      # When we don't have a static batch size.
+      raise ValueError('DTensor strategy requires a static batch size for now.'
+                       'The dynamic batch size will be supported in future')
+    global_batch_size = batch_size.numpy() * self._num_replicas_in_sync
+
+    return input_util.DTensorDataset(
+        dataset=dataset,
+        mesh=self._mesh,
+        layouts=layouts,
+        global_batch_size=global_batch_size,
+        dataset_already_batched=True,
+        batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+        # TODO(scottzhu): Add prefetch support by inspecting the input dataset.
+        prefetch=None,
+        tf_data_service_config=None
+    )
+
+  def _experimental_distribute_values_from_function(self, value_fn):
+    per_replica_values = []
+    # Note that in the multi-worker setting, this function only return the
+    # slide of DistributedValue for the current worker.
+    for i in range(self._mesh.num_local_devices()):
+      # In the case of 2 worker with 2 local devices on each worker,
+      # worker 0 will get 0 and 1 for replica_id.
+      # worker 1 will get 2 and 3 for replica_id.
+      replica_id = d_config.client_id() * self._mesh.num_local_devices() + i
+      per_replica_values.append(value_fn(
+          distribute_lib.ValueContext(replica_id,
+                                      self._num_replicas_in_sync)))
+    # Instead of using the DistributeVariable, return a DTensor instead since
+    # the run() will expect a DTensor instance.
+    result = distribute_utils.regroup(per_replica_values, always_wrap=True)
+    map_fn = functools.partial(dtensor_util.convert_per_replica_to_dtensor,
+                               mesh=self._mesh)
+    return nest.map_structure(map_fn, result)
+
+  def call_for_each_replica(self, fn, args=(), kwargs=None):
+    """Run `fn` once per replica.
+
+    This is a method that expected by the strategy base class in its `run()`.
+
+    Args:
+      fn: function to run (will be run once per replica).
+      args: Tuple or list with positional arguments for `fn`.
+      kwargs: Dict with keyword arguments for `fn`.
+
+    Returns:
+      Merged return value of `fn` across all replicas.
+    """
+    # Comparing to the existing MirroredStrategy, which will run the fn on
+    # each of the replica with individual thread, the DTensor will just run
+    # the fn once with the DTensor inputs, and the distribution will be handled
+    # by the DTensor.
+
+    distribute_lib._require_cross_replica_or_default_context_extended(self)   # pylint: disable=protected-access
+    if kwargs is None:
+      kwargs = {}
+
+    # For any value that is not DTensor, eg normal tf.Tensor or
+    # DistributedValues, we need to convert them into DTensor.
+    map_fn = functools.partial(dtensor_util.convert_inputs_to_dtensor,
+                               mesh=self._mesh)
+    d_args = nest.map_structure(map_fn, args)
+    d_kwargs = nest.map_structure(map_fn, kwargs)
+
+    with self._container_strategy().scope():
+      with dtensor_util.DTensorReplicaContext(self._container_strategy()):
+        dtensor_result = fn(*d_args, **d_kwargs)
+
+    return nest.map_structure(
+        dtensor_util.DTensorDistributedValue,
+        dtensor_result)
+
+  def _gather_to_implementation(self, value, destinations, axis, options):
+    if isinstance(value, dtensor_util.DTensorDistributedValue):
+      value = value.get_dtensor()
+    if not d_api.is_dtensor(value):
+      # This is the current behavior for mirrored strategy, should we raise an
+      # error for unsupported types?
+      return value
+
+    # Unpack the dtensor components and gather the tensors on the axis
+    components = d_api.unpack(value)
+    return array_ops.concat(components, axis=axis)
+
+  def _use_merge_call(self):
+    # This is method for V1 StrategyExtended by still used by
+    # tf.__internal__.distribute.strategy_supports_no_merge_call
+    return False

tensorflow/python/distribute/experimental/dtensor_util.py

@@ -0,0 +1,299 @@
+# 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.
+# ==============================================================================
+"""Utilities for strategies that are backed by DTensor."""
+
+from tensorflow.dtensor.python import accelerator_util
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import input_util
+from tensorflow.dtensor.python import layout
+from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion_registry
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import summary_ops_v2
+
+
+# Default dimension name used for the mesh created when user provide a list
+# of devices. For mirrored strategy, it should be a 1D mesh with batch dim only.
+DEFAULT_BATCH_MESH_DIM_NAME = "batch"
+
+
+class DTensorDistributedValue(values_lib.DistributedValues):
+  """DistributedValue backed by a DTensor instance.
+
+  This class is useful to align the interface between DTensor and tf.distribute.
+  Most of the tf.distribute API will accept/return DistributedValue, whereas
+  DTensor low level API will only accept DTensor instance. In order to avoid
+  the conversion back and forth between DistributedValue and DTensor, we
+  introduce this class so that it can work with both side.
+  """
+
+  def __init__(self, dtensor):
+    if context.executing_eagerly():
+      if not d_api.is_dtensor(dtensor):
+        raise ValueError("The DTensorDistributedValue can only be built with "
+                         f"DTensor instance, got {type(dtensor)}")
+      super().__init__(d_api.unpack(dtensor))
+    else:
+      # We can't unpack the dtensor instance for now due to graph context.
+      # We will treat the dtensor instance as one global instance and let it
+      # return as a global replica instance.
+      # TODO(feyu): Support unpack in the graph context.
+      super().__init__([dtensor,])
+    self._dtensor = dtensor
+
+  def get_dtensor(self):
+    return self._dtensor
+
+  @property
+  def values(self):
+    # Note that this method exists so that it match the interface for PerReplica
+    # The public API in `tf.types.experimental.distributed.PerReplica` doesn't
+    # define any methods.
+    return self._values
+
+
+def _dtensor_distributed_value_to_tensor(
+    var, dtype=None, name=None, as_ref=False):
+  del name
+  dtensor = var.get_dtensor()
+  if dtype is not None and not dtype.is_compatible_with(dtensor.dtype):
+    raise ValueError(
+        "Incompatible type conversion requested to type {!r} for variable "
+        "of type {!r}".format(dtype.name, dtensor.dtype.name))
+  if as_ref:
+    raise NotImplementedError(
+        "PerReplica doesn't support being used as a reference.")
+  return dtensor
+
+
+# Register a conversion function to provide a useful error message when users
+# try to use PerReplica values in the wrong contexts
+tensor_conversion_registry.register_tensor_conversion_function(
+    DTensorDistributedValue, _dtensor_distributed_value_to_tensor)
+
+
+class DTensorReplicaContext(distribute_lib.ReplicaContext):
+  """ReplicaContext for strategy that is backed by DTensor.
+
+  Since the DTensor is operated in the global context, most of the methods from
+  existing strategy ReplicaContext is not applicable since they need to access
+  local values. For now most of the methods in this class will raise explicit
+  error to user, and we will add more support for local values in future.
+  """
+  _UNSUPPORTED_ERROR_MSG = (
+      "Strategy that is backed by DTensor is run with a global context, and "
+      "doesn't support operations for local context, like any call to merge/"
+      "gather/reduce or local replica ID. Please use any strategy that is not "
+      "backed by DTensor")
+
+  def __init__(self, strategy):
+    # Since DTensor strategy only runs in a global context, and we can't have
+    # a local replica ID in the sync group. For now we pass None to parent, and
+    # raise an explicit error when it is accessed.
+    super().__init__(strategy, replica_id_in_sync_group=None)
+
+  def __enter__(self):
+    # This is a copy of parent class, without any check about whether the
+    # current replica is the first one (since DTensor only has one).
+    distribute_lib._push_per_thread_mode(self._thread_context)  # # pylint: disable=protected-access
+
+    summary_state = summary_ops_v2._summary_state  # pylint: disable=protected-access
+    self._summary_recording_distribution_strategy = (
+        summary_state.is_recording_distribution_strategy)
+    summary_state.is_recording_distribution_strategy = True
+
+  @property
+  def replica_id_in_sync_group(self):
+    # Since there is only one global context for DTensor, we always return a
+    # constant value here. This value is needed by the RNG which try to generate
+    # different seed for different replica.
+    return 0
+
+  @property
+  def _replica_id(self):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def merge_call(self, merge_fn, args=(), kwargs=None):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def all_reduce(self, reduce_op, value, options=None):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def all_gather(self, value, axis, options=None):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def _update(self, var, fn, args=(), kwargs=None, group=True):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+
+def initialize_accelerator_system_once(device_type):
+  # Initialize the GPU/TPU before creating the mesh.
+  # Note that this method will also trigger the creation of the pairing
+  # virtual host CPUs, which is needed by dataset and checkpoint.
+  if not accelerator_util.is_initialized():
+    # TODO(feyu): Add a method in accelerator_util to check the initialized
+    # mesh device types.
+    accelerator_util.initialize_accelerator_system(
+        device_type,
+        experimental_reset_context=True)
+
+
+def convert_inputs_to_dtensor(inputs, mesh):
+  """Convert any input types to DTensor instance."""
+  if isinstance(inputs, DTensorDistributedValue):
+    return inputs.get_dtensor()
+  elif isinstance(inputs, values_lib.DistributedValues):
+    return convert_per_replica_to_dtensor(inputs, mesh)
+  elif isinstance(inputs, input_util._DTensorIterator):   # pylint: disable=protected-access
+    return inputs
+  elif tensor_util.is_tensor(inputs):
+    if context.executing_eagerly():
+      if d_api.is_dtensor(inputs):
+        return inputs
+      else:
+        # For a non-dtensor input in eager context, we could choose to replica
+        # them into per-replica and then pack them into dtensor. However, this
+        # will cause an eager/graph discrepancy since we can't do this check in
+        # the graph context. For now, we will ask user to provide a distributed
+        # value for inputs.
+        _raise_unsupported_input_type_error(inputs)
+    else:
+      # For graph context, since we can't check if they are dtensor or not. We
+      # will assume the value is already distributed. This is a critical use
+      # case for keras, where all the inputs are pre-distributed via strategy,
+      # and the train function execute within graph context.
+      return inputs
+  else:
+    # For any other types.
+    _raise_unsupported_input_type_error(inputs)
+
+
+def _raise_unsupported_input_type_error(inputs):
+  raise ValueError("Unsupported input types for MirroredStrategy. "
+                   "Please use `strategy.distribute_dataset` or "
+                   "`strategy.distribute_values_from_function` to "
+                   f"distribute inputs. Received input type: {type(inputs)}")
+
+
+def is_distributed_value(value):
+  return isinstance(
+      value, values_lib.DistributedValues) or d_api.is_dtensor(value)
+
+
+def convert_per_replica_to_dtensor(per_replica_value, mesh):
+  """Convert a PerReplica result to a DTensor instance.
+
+  Args:
+    per_replica_value: A PerReplica instance whose value will be converted
+      to DTensor.
+    mesh: The mesh used for layout creation.
+
+  Returns:
+    A DTensor instance that packed from per_replica_value with batch sharded
+      layout.
+  """
+  values = per_replica_value.values
+  if isinstance(values[0], (float, int)):
+    rank = 0
+  else:
+    rank = len(values[0].shape)
+
+  if rank == 0:
+    result = []
+    # dtensor.pack requires each component to have same rank as the packed
+    # result. When the individual value is scalar, it needs to be expanded into
+    # 1D tensor.
+    for v in values:
+      result.append(array_ops.expand_dims_v2(v, axis=0))
+    rank += 1
+  else:
+    result = list(values)   # dtensor.pack requires a list as input.
+
+  # TODO(scottzhu): Note that the result tensor could be a partial value and
+  # not always batch shard or fully replicaed. See
+  # http://screenshot/6ERkXyX95KqftCw as an example.
+  batch_layout = layout.Layout.batch_sharded(
+      mesh, batch_dim=DEFAULT_BATCH_MESH_DIM_NAME, rank=rank)
+
+  return d_api.pack(result, batch_layout)
+
+
+def dtensor_reduce(strategy, reduce_op, value, axis):
+  """Implement dtensor based strategy.reduce()."""
+  # Due to the limitation of using scalar in DTensor (e.g. the rank 0 tensor
+  # loss the batch shard information), we need to override the default
+  # reduce in addition to the strategy.extend._reduce_to()
+  # Most of the logic here is a mimic of the parent class, except for how
+  # mean and sum are calculated in a global context.
+  distribute_lib._require_cross_replica_or_default_context_extended(  # pylint: disable=protected-access
+      strategy.extended)
+  if isinstance(reduce_op, str):
+    reduce_op = reduce_util.ReduceOp(reduce_op.upper())
+
+  distributed_input = is_distributed_value(value)
+  if not distributed_input and axis is None:
+    # For any value that isn't distributed and doesn't need a reduction within
+    # the replica.
+    destinations = (device_util.current() or
+                    strategy.extended._default_device or  # pylint: disable=protected-access
+                    "/device:CPU:0")
+    devices = cross_device_ops_lib.get_devices_from(destinations)
+    with ops.device(devices[0]):
+      return array_ops.identity(
+          cross_device_ops_lib.reduce_non_distributed_value(
+              reduce_op, value, destinations, strategy.num_replicas_in_sync))
+
+  value = convert_inputs_to_dtensor(value, strategy._mesh)  # pylint: disable=protected-access
+  # At this point, the value is a DTensor instance now.
+  # There will be a final reduction step cross replica. In order to maintain
+  # the shape of each local replica, we need to add a new dim to the front.
+  # E.g. 2 replica with local shape as (4, 5, 6), the global tensor shape
+  # should be (8, 5, 6), we will reshape into (2, 4, 5, 6) and then do a
+  # reduction on axis 0.
+  if reduce_op == reduce_util.ReduceOp.MEAN:
+    reduce_op = math_ops.reduce_mean
+  else:
+    reduce_op = math_ops.reduce_sum
+
+  # TODO(scottzhu): Make sure we handle dynamic/uneven shape in future.
+  if d_api.fetch_layout(value).is_fully_replicated():
+    # In case of fully mirrored dtensor, we only need to do one reduce, and
+    # don't need to care about any per-replica logic.
+    if axis is not None:
+      value = reduce_op(value, axis=axis)
+  else:
+    new_shape = [strategy.num_replicas_in_sync, -1]
+    if len(value.shape) > 1:
+      new_shape.extend(array_ops.shape(value)[1:])
+    value = array_ops.reshape(value, new_shape)
+    if axis is not None:
+      # we do a reduce_sum/mean within each of the replica when axis is not
+      # None. Add 1 to the axis since there is a new dim added by reshape in
+      # front.
+      value = reduce_op(value, axis=axis + 1)
+    value = reduce_op(value, axis=0)
+
+  # Note that we return a DTensor instance here, which should have the same
+  # value as the original MirroredStrategy, but with a different type. User
+  # might want a tf.Tensor for the status quo.
+  return value

tensorflow/python/distribute/experimental/dtensor_util_test.py

@@ -0,0 +1,122 @@
+# 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.
+# ==============================================================================
+"""Test for DTensor related utilities in tf.distribute."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.distribute.experimental import mirrored_strategy
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+
+
+class DTensorDistributedValueTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+    tensor_1 = constant_op.constant([1.0])
+    tensor_2 = constant_op.constant([2.0])
+    self.batch_layout = layout.Layout.batch_sharded(
+        self.mesh, batch_dim='batch', rank=1)
+    self.dtensor = d_api.pack([tensor_1, tensor_2], self.batch_layout)
+
+  @parameterized.named_parameters([
+      ('py_floats', [1.0, 2.0]),
+      ('np_floats', np.array([1.0, 2.0])),
+      ('tf_const', lambda: constant_op.constant([1.0, 2.0])),
+      ('distribute_value', values_lib.PerReplica([1.0, 2.0])),
+  ])
+  def test_input_validation(self, inputs):
+    if callable(inputs):
+      inputs = inputs()
+
+    with self.assertRaisesRegex(ValueError, 'can only be built with DTensor'):
+      dtensor_util.DTensorDistributedValue(inputs)
+
+  def test_unpack(self):
+    v = dtensor_util.DTensorDistributedValue(self.dtensor)
+
+    self.assertIs(self.dtensor, v.get_dtensor())
+
+    per_replica_result = v.values
+    self.assertLen(per_replica_result, 2)
+    self.assertAllClose(per_replica_result[0], constant_op.constant([1.0]))
+    self.assertAllClose(per_replica_result[1], constant_op.constant([2.0]))
+
+  def test_graph_behavior(self):
+
+    @def_function.function
+    def run_fn(input_dtensor):
+      return dtensor_util.DTensorDistributedValue(input_dtensor)
+
+    result = run_fn(self.dtensor)
+    # When it cross the boundary of tf.function, it will be unwrapped and
+    # return a dtensor instance directly.
+    self.assertTrue(d_api.is_dtensor(result))
+    self.assertDTensorEqual(constant_op.constant([1.0, 2.0]),
+                            self.batch_layout, result)
+
+
+class DTensorReplicaContextTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+  def test_unsupported_methods(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    replica_context = dtensor_util.DTensorReplicaContext(strategy)
+
+    expected_error = replica_context._UNSUPPORTED_ERROR_MSG
+
+    self.assertEqual(replica_context.num_replicas_in_sync, 2)
+
+    self.assertEqual(replica_context.replica_id_in_sync_group, 0)
+
+    with self.assertRaisesRegex(NotImplementedError, expected_error):
+      replica_context.merge_call(None)
+
+    with self.assertRaisesRegex(NotImplementedError, expected_error):
+      replica_context.all_reduce(reduce_util.ReduceOp.SUM, None)
+
+    with self.assertRaisesRegex(NotImplementedError, expected_error):
+      replica_context.all_gather([], 0)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/experimental/mirrored_strategy.py

@@ -0,0 +1,99 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Implement a MirroredStrategy based on the DTensor low level API.
+
+This is an experiment to validate the viability of the DTensor API, and expose
+any potential feature gaps between the current API and the need.
+"""
+
+from tensorflow.dtensor.python import config as d_config
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute.experimental import dtensor_strategy_extended
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.framework import device as tf_device
+
+
+class MirroredStrategy(distribute_lib.Strategy):
+  """Synchronous training across multiple replicas on one machine.
+
+  This strategy is typically used for training on one machine with multiple
+  accelerators (GPUs/TPUs).
+
+  For example, a variable created under a `MirroredStrategy` is a distributed
+  variable with layout replicated on each dimension. The variables will be
+  placed on the `mesh` that is specified in the __init__.
+  """
+
+  def __init__(self, devices=None, cross_device_ops=None, *, mesh=None):
+    """Synchronous training across multiple replicas on one machine.
+
+    Args:
+      devices: a list of device strings, such as ['/gpu:0', '/gpu:1']. If both
+        `mesh` and `devices` are None, all the available GPU/TPU will be used.
+        If no accelerators are found, CPU is used.
+      cross_device_ops: optional, a descendant of `CrossDeviceOps`. The value is
+        ignored at the moment, and support will be added later.
+      mesh: optional DTensor mesh for the computation. Note that either `mesh`
+        or `devices` should be provided, and not both. The mesh should be 1D,
+        and will be used to split the input data among that dimension.
+    """
+    self._validate_init_args(mesh, devices)
+    if not mesh:
+      mesh = self._build_mesh_from_device_list(devices)
+
+    extended = dtensor_strategy_extended.DTensorStrategyExtended(
+        container_strategy=self, mesh=mesh)
+    super().__init__(extended)
+    self._mesh = mesh
+    self._devices = devices
+
+  @classmethod
+  def _validate_init_args(cls, mesh, devices):
+    if mesh and devices:
+      raise ValueError('Mesh and devices can not be provided at the same time. '
+                       f'received mesh = {mesh}, devices = {devices}')
+
+    # For mirrored strategy, the mesh should be 1D, and only contains a batch
+    # dimension, we will use that dimension to shard the inputs.
+    if mesh and len(mesh.shape()) != 1:
+      raise ValueError('The mesh for MirroredStrategy must be 1D, received: '
+                       f'{len(mesh.shape())}D')
+
+  @classmethod
+  def _build_mesh_from_device_list(cls, devices):
+    if devices:
+      device_type = tf_device.DeviceSpec.from_string(devices[0]).device_type
+      dtensor_util.initialize_accelerator_system_once(device_type)
+      mesh = mesh_util.create_mesh(
+          mesh_dims=[(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME, len(devices))],
+          devices=devices)
+    else:
+      # Trying to detect if there is any GPU/TPUs attached.
+      device_type = d_config.preferred_device_type()
+      devices = d_config.local_devices(device_type)
+      dtensor_util.initialize_accelerator_system_once(device_type)
+      mesh = mesh_util.create_mesh(
+          mesh_dims=[(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME, len(devices))],
+          device_type=device_type)
+    return mesh
+
+  def reduce(self, reduce_op, value, axis):
+    return dtensor_util.dtensor_reduce(self, reduce_op, value, axis)
+
+  @property
+  def mesh(self):
+    """Returns the mesh used by the strategy."""
+    return self._mesh

tensorflow/python/distribute/experimental/mirrored_strategy_test.py

@@ -0,0 +1,722 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test for MirroredStrategy backed by DTensor API."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import d_variable
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.distribute.experimental import mirrored_strategy
+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 dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.ops import variables
+
+
+class StrategyBaseTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+  @parameterized.named_parameters([
+      ('py_floats', lambda: [1.0, 2.0], True),
+      ('np_floats', lambda: np.array([1.0, 2.0]), True),
+      ('tf_const', lambda: constant_op.constant([1.0, 2.0]), True),
+      ('py_floats_callable', lambda: [1.0, 2.0], False),
+      ('np_floats_callable', lambda: np.array([1.0, 2.0]), False),
+      ('tf_const_callable', lambda: constant_op.constant([1.0, 2.0]), False),
+  ])
+  def test_variable_creation(self, init_value, convert_callable):
+    if convert_callable:
+      init_value = init_value()
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy.scope():
+      v = variables.Variable(init_value)
+
+    self.assertIsInstance(v, d_variable.DVariable)
+    self.assertIsNotNone(v.layout)
+    self.assertEqual(v.layout, layout.Layout.replicated(self.mesh, rank=1))
+
+  def test_variable_creation_with_dtype(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy.scope():
+      v = variables.Variable(
+          0, dtype='int64',
+          aggregation=variables.VariableAggregationV2.ONLY_FIRST_REPLICA)
+    self.assertIsInstance(v, d_variable.DVariable)
+    self.assertEqual(v.dtype, dtypes.int64)
+
+  def test_mesh(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertEqual(strategy.mesh, self.mesh)
+
+  def test_strategy_extension(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertIsInstance(strategy.extended, distribute_lib.StrategyExtendedV2)
+
+  def test_num_replica_in_sync(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertEqual(strategy.num_replicas_in_sync, 2)
+
+  def test_worker_devices(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    worker_devices = strategy.extended.worker_devices
+    self.assertLen(worker_devices, 2)
+    self.assertEqual(worker_devices, tuple(self.mesh.local_devices()))
+
+  def test_parameter_devices(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    parameter_devices = strategy.extended.parameter_devices
+    self.assertLen(parameter_devices, 2)
+    self.assertEqual(parameter_devices, tuple(self.mesh.local_devices()))
+
+  def test_variable_created_in_scope(self):
+    strategy1 = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy1.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    v2 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    strategy2 = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy2.scope():
+      v3 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    self.assertTrue(strategy1.extended.variable_created_in_scope(v1))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v2))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v3))
+    self.assertTrue(strategy2.extended.variable_created_in_scope(v3))
+
+  def test_colocate_vars_with(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+      with strategy.extended.colocate_vars_with(v1):
+        v2 = variables.Variable(constant_op.constant([2.0, 3.0]))
+
+    # We assert the layout for the variable, and make sure they are same.
+    self.assertEqual(v1.layout, v2.layout)
+
+  def test_in_multi_worker_mode(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertFalse(strategy.extended._in_multi_worker_mode())
+
+  def test_run_with_tensor_inputs(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant(3.0)
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2.0
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.run(replica_fn, args=(tensor_input,))
+
+  def test_run_with_graph_tensor_inputs(self):
+    # Note that this is potentially a sharp edge for the user, since the eager
+    # test case was raising an error, but the graph context will run, by treat
+    # the inputs as a global inputs.
+    # TODO(scottzhu): Mitigate this eager/graph behavior difference in future.
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2.0
+
+    @def_function.function
+    def run_fn():
+      tensor_input = constant_op.constant(3.0)
+      return strategy.run(replica_fn, args=(tensor_input,))
+
+    with strategy.scope():
+      result = run_fn()
+    self.assertEqual(result, constant_op.constant(6.0))
+
+  def test_run_with_unsupported_input_types(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    random_inputs = [123, '456']
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2.0
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.run(replica_fn, args=(random_inputs,))
+
+  def test_run_with_distribute_value_input(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    def value_fn(value_context):
+      return value_context.replica_id_in_sync_group
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, 2)
+    # Note that the scalar value from
+    # experimental_distribute_values_from_function will be up rank to 1D since
+    # batched shared dtensor need at least be 1D. So the result from the
+    # strategy.run is [0], instead of just 0.
+    self.assertAllClose(result.values[0], constant_op.constant([0]))
+    self.assertAllClose(result.values[1], constant_op.constant([2]))
+
+  def test_run_without_input(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([1.0])
+
+    result = strategy.run(replica_fn)
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, 2)
+
+    self.assertAllClose(result.values[0], constant_op.constant([1.0]))
+    self.assertAllClose(result.values[1], constant_op.constant([1.0]))
+
+  def test_nested_structure_output(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    array_value = np.array([3., 2., 1.])
+    def value_fn(ctx):
+      value = array_value[ctx.replica_id_in_sync_group]
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertLen(result.keys(), 2)
+    self.assertIsInstance(result['a'], dtensor_util.DTensorDistributedValue)
+    self.assertAllClose(result['a'].values[0], constant_op.constant([6.0]))
+    self.assertAllClose(result['a'].values[1], constant_op.constant([4.0]))
+
+    self.assertIsInstance(result['b'], dtensor_util.DTensorDistributedValue)
+    self.assertAllClose(result['b'].values[0],
+                        constant_op.constant([8.0, 10.0]))
+    self.assertAllClose(result['b'].values[1], constant_op.constant([6.0, 8.0]))
+
+  def test_inputs_with_dtensor_distribute_values(self):
+
+    @def_function.function
+    def replica_fn_1(inputs):
+      return inputs * 2.0
+
+    @def_function.function
+    def replica_fn_2(inputs):
+      return inputs + 1.0
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant(3.0)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    result_1 = strategy.run(replica_fn_1, args=(d_tensor_input,))
+    self.assertIsInstance(result_1, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_1.values, 2)
+    self.assertAllClose(result_1.values[0], constant_op.constant([6.0]))
+    self.assertAllClose(result_1.values[1], constant_op.constant([6.0]))
+
+    result_2 = strategy.run(replica_fn_2, args=(result_1,))
+    self.assertIsInstance(result_2, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_2.values, 2)
+    self.assertAllClose(result_2.values[0], constant_op.constant([7.0]))
+    self.assertAllClose(result_2.values[1], constant_op.constant([7.0]))
+
+  def test_run_with_nullary_ops(self):
+
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([3.0])
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    result = strategy.run(replica_fn)
+
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertAllClose(result.values[0], constant_op.constant([3.0]))
+    self.assertAllClose(result.values[1], constant_op.constant([3.0]))
+
+  def test_get_replica_context(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    tensor_input = constant_op.constant(3)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    @def_function.function
+    def replica_fn(inputs):
+      replica_context = distribute_lib.get_replica_context()
+      self.assertIsInstance(replica_context, dtensor_util.DTensorReplicaContext)
+      return inputs * replica_context.num_replicas_in_sync
+
+    # Default replica context
+    self.assertIsNotNone(distribute_lib.get_replica_context())
+    with strategy.scope():
+      self.assertIsNone(distribute_lib.get_replica_context())
+
+      result = strategy.run(replica_fn, args=(d_tensor_input,))
+
+    self.assertLen(result.values, 2)
+    self.assertAllClose(result.values[0], constant_op.constant([6]))
+    self.assertAllClose(result.values[1], constant_op.constant([6]))
+
+  def test_gather_non_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant(3.0)
+
+    result = strategy.gather(tensor_input, axis=0)
+    self.assertAllClose(result, tensor_input)
+
+  def test_gather_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    def value_fn(value_context):
+      start = value_context.replica_id_in_sync_group
+      return array_ops.reshape(math_ops.range(start=start, limit=start + 6),
+                               shape=(1, 2, 3))
+
+    distribute_result = strategy.experimental_distribute_values_from_function(
+        value_fn)
+    result = strategy.gather(distribute_result, axis=0)
+    self.assertEqual(result.shape, [2, 2, 3])
+    self.assertAllClose(result, [[[0, 1, 2], [3, 4, 5]],
+                                 [[1, 2, 3], [4, 5, 6]]])
+    result = strategy.gather(distribute_result, axis=1)
+    self.assertEqual(result.shape, [1, 4, 3])
+    self.assertAllClose(result, [[[0, 1, 2], [3, 4, 5], [1, 2, 3], [4, 5, 6]]])
+    result = strategy.gather(distribute_result, axis=2)
+    self.assertEqual(result.shape, [1, 2, 6])
+    self.assertAllClose(result, [[[0, 1, 2, 1, 2, 3], [3, 4, 5, 4, 5, 6]]])
+
+  def test_reduce_mean_non_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=0)
+
+  def test_reduce_sum_non_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.reduce(reduce_util.ReduceOp.SUM, tensor_input, axis=0)
+
+  def test_reduce_mean_distribute_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      n = value_context.num_replicas_in_sync
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * n * 2.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn)
+    # replica 1 has [[0.0, 1.0],[2.0, 3.0]] and replica 2 has
+    # [[4.0, 5.0],[6.0, 7.0]]
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=None)
+    self.assertAllClose(result, constant_op.constant([[2.0, 3.0], [4.0, 5.0]]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=0)
+    self.assertAllClose(result, constant_op.constant([3.0, 4.0]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=1)
+    self.assertAllClose(result, constant_op.constant([2.5, 4.5]))
+
+  def test_reduce_sum_distribute_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      n = value_context.num_replicas_in_sync
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * n * 2.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn)
+    # replica 1 has [[0.0, 1.0],[2.0, 3.0]] and replica 2 has
+    # [[4.0, 5.0],[6.0, 7.0]]
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=None)
+    self.assertAllClose(result, constant_op.constant([[4.0, 6.0], [8.0, 10.0]]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=0)
+    self.assertAllClose(result, constant_op.constant([12.0, 16.0]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=1)
+    self.assertAllClose(result, constant_op.constant([10.0, 18.0]))
+
+  def test_reduce_mean_mirrored_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    with  strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([2.0, 3.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([1.5, 3.5]))
+
+  def test_reduce_sum_mirrored_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    with  strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([4.0, 6.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([3.0, 7.0]))
+
+  def test_reduce_value_device(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=None)
+    self.assertIn('CPU:0', result.device)
+
+  def test_experimental_local_results(self):
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([3.0])
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    result = strategy.run(replica_fn)
+    local_result = strategy.experimental_local_results(result)
+
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, 2)
+    self.assertEqual(local_result[0], constant_op.constant([3.0]))
+    self.assertEqual(local_result[1], constant_op.constant([3.0]))
+
+  def test_experimental_local_results_with_inputs(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    array_value = np.array([3., 2.])
+    def value_fn(ctx):
+      value = array_value[ctx.replica_id_in_sync_group]
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    local_result = strategy.experimental_local_results(result)
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, 2)
+    self.assertDictEqual(local_result[0],
+                         {'a': constant_op.constant([6.0]),
+                          'b': constant_op.constant([8.0, 10.0])})
+    self.assertDictEqual(local_result[1],
+                         {'a': constant_op.constant([4.0]),
+                          'b': constant_op.constant([6.0, 8.0])})
+
+
+class InvalidMeshTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2, 1))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch', 'model'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh_2d = self.configTestMesh(mesh_dict)
+
+  def test_invalid_mesh_shape(self):
+    with self.assertRaisesRegex(
+        ValueError, 'The mesh for MirroredStrategy must be 1D, received: 2D'):
+      mirrored_strategy.MirroredStrategy(mesh=self.mesh_2d)
+
+
+class StrategyCreationTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    device_type = test_util.preferred_device_type()
+    if device_type != 'TPU':
+      test_util.reset_logical_devices(device_type, 2)
+    self.device_type = device_type
+
+  def test_explicit_device_list(self):
+    device_list = [f'/{self.device_type}:{i}' for i in range(2)]
+    strategy = mirrored_strategy.MirroredStrategy(devices=device_list)
+    mesh = strategy.mesh
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.shape(), [2,])
+    self.assertEqual(mesh.dim_names, ['batch'])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:0',
+        mesh.local_devices()[0])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:1',
+        mesh.local_devices()[1])
+    # Also make sure the host mesh works since it is required by dataset
+    self.assertIsNotNone(mesh.host_mesh())
+
+  def test_implicit_device_list(self):
+    strategy = mirrored_strategy.MirroredStrategy()
+    mesh = strategy.mesh
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.shape(), [2,])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:0',
+        mesh.local_devices()[0])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:1',
+        mesh.local_devices()[1])
+    # Also make sure the host mesh works since it is required by dataset
+    self.assertIsNotNone(mesh.host_mesh())
+
+  def test_mesh_with_device_list(self):
+    device_list = [f'/{self.device_type}:{i}' for i in range(2)]
+    mesh = mesh_util.create_mesh([('batch', 2)], devices=device_list)
+    with self.assertRaisesRegex(
+        ValueError, 'Mesh and devices can not be provided at the same time'):
+      _ = mirrored_strategy.MirroredStrategy(mesh=mesh, devices=device_list)
+
+
+class StrategyDatasetTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+    self.images = stateless_random_ops.stateless_random_uniform(
+        [8, 8, 3], seed=(1, 2), minval=0, maxval=255)
+    self.labels = stateless_random_ops.stateless_random_uniform(
+        [1], seed=(1, 2), minval=0, maxval=10)
+
+    self.dataset = dataset_ops.Dataset.from_tensors(
+        (self.images, self.labels)).repeat()
+
+  def test_create_batched_dataset(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    global_batch_size = 8
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    element = next(iter(distributed_dataset))
+    batched_image, batched_label = element
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # Make sure when unpack the tensor, each of them has enough shards.
+    self.assertLen(d_api.unpack(batched_image), self.mesh.num_local_devices())
+    self.assertLen(d_api.unpack(batched_label), self.mesh.num_local_devices())
+
+  def test_uneven_batched_dataset(self):
+    elements = [[1, 2, 3], [1, 2], [1, 2, 3, 4]]
+    dataset = dataset_ops.Dataset.from_generator(
+        lambda: elements, dtypes.int64).repeat()
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with self.assertRaisesRegex(ValueError, 'requires a static batch size'):
+      strategy.experimental_distribute_dataset(dataset)
+
+  def test_create_partial_batched_dataset(self):
+    # TODO(b/210887657): Support last partial batch.
+    self.skipTest('Test failed due to last partial batch')
+    dataset = dataset_ops.Dataset.from_tensors(
+        (self.images, self.labels)).repeat(30)  # There is a last partial batch
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    global_batch_size = 8
+    dataset = dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    expected_element_batch_size = [8, 8, 8, 6]
+    # The last batch with 6 element will fail to produce with StopIteration.
+    iterator = iter(distributed_dataset)
+    for batch_size in expected_element_batch_size:
+      element = next(iterator)
+      batched_image, batched_label = element
+      self.assertEqual(batched_image.shape, [batch_size, 8, 8, 3])
+      self.assertEqual(batched_label.shape, [batch_size, 1])
+
+      # Make sure when unpack the tensor, each of them has enough shards.
+      self.assertLen(d_api.unpack(batched_image), self.mesh.num_local_devices())
+      self.assertLen(d_api.unpack(batched_label), self.mesh.num_local_devices())
+
+  def test_deprecated_strategy_methods(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_dataset_iterator(self.dataset)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_input_fn_iterator(lambda _: self.dataset)
+
+  def test_distribute_dataset_from_fn(self):
+    local_batch_size = 4
+    global_batch_size = 8
+    def dataset_fn(option):
+      del option
+      return dataset_ops.Dataset.from_tensors(
+          (self.images, self.labels)).repeat().batch(
+              local_batch_size, drop_remainder=True).prefetch(2)
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    distributed_dataset = strategy.distribute_datasets_from_function(
+        dataset_fn, None)
+    iterator = iter(distributed_dataset)
+
+    self.assertEqual(distributed_dataset.element_spec,
+                     (tensor_spec.TensorSpec(shape=(8, 8, 8, 3),
+                                             dtype=dtypes.float32, name=None),
+                      tensor_spec.TensorSpec(shape=(8, 1),
+                                             dtype=dtypes.float32, name=None)))
+    self.assertEqual(distributed_dataset.element_spec, iterator.element_spec)
+
+    batched_image, batched_label = next(iterator)
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # Make sure there are two shards when unpack, and each of them has 4 as
+    # batch size
+    unpacked_images = d_api.unpack(batched_image)
+    self.assertLen(unpacked_images, self.mesh.num_local_devices())
+    self.assertEqual(unpacked_images[0].shape, [local_batch_size, 8, 8, 3])
+    self.assertEqual(unpacked_images[1].shape, [local_batch_size, 8, 8, 3])
+
+  def test_distribute_values_from_function(self):
+    array_value = np.array([3., 2., 1.])
+    def value_fn(ctx):
+      return array_value[ctx.replica_id_in_sync_group]
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+    self.assertDTensorEqual(
+        constant_op.constant([3., 2.], dtype=dtypes.float64),
+        layout.Layout.batch_sharded(self.mesh, batch_dim='batch', rank=1),
+        distributed_values)
+
+  def test_distribute_values_from_function_with_nested_structure(self):
+    array_value = np.array([3., 2., 1.])
+    def value_fn(ctx):
+      value = array_value[ctx.replica_id_in_sync_group]
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+    self.assertIsInstance(distributed_values, dict)
+    self.assertDTensorEqual(
+        constant_op.constant([3., 2.], dtype=dtypes.float64),
+        layout.Layout.batch_sharded(self.mesh, batch_dim='batch', rank=1),
+        distributed_values['a'])
+    unpacked_a = d_api.unpack(distributed_values['a'])
+    # Note that this might have a slight behavior difference, the original
+    # mirrored strategy may return scalar for each PerReplica. The DTensor
+    # implementation is more strict and always ensures the PerReplica
+    # value has the same rank as the global-view Tensor.
+    self.assertAllClose(unpacked_a[0], [3.])
+    self.assertAllClose(unpacked_a[1], [2.])
+    self.assertDTensorEqual(
+        constant_op.constant([4., 5., 3., 4.], dtype=dtypes.float64),
+        layout.Layout.batch_sharded(self.mesh, batch_dim='batch', rank=1),
+        distributed_values['b'])
+
+  def test_distribute_dataset_in_tf_function(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    local_batch_size = 4
+    global_batch_size = 8
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+
+    @def_function.function
+    def step_fn(iterator):
+      images, labels = next(iterator)
+      del labels
+      return images
+
+    result = strategy.run(step_fn, args=(iter(distributed_dataset),))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, self.mesh.num_local_devices())
+    self.assertEqual(result.values[0].shape, [local_batch_size, 8, 8, 3])
+    self.assertEqual(result.values[1].shape, [local_batch_size, 8, 8, 3])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/experimental/multi_worker_mirrored_strategy.py

@@ -0,0 +1,156 @@
+# 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.
+# ==============================================================================
+"""Implement a MultiMirroredStrategy based on the DTensor low level API.
+
+This is an experiment to validate the viability of the DTensor API, and expose
+any potential feature gaps between the current API and the need.
+"""
+
+import os
+
+from tensorflow.dtensor.python import config as d_config
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.distribute.experimental import dtensor_strategy_extended
+from tensorflow.python.distribute.experimental import dtensor_util
+
+
+class MultiWorkerMirroredStrategy(distribute_lib.Strategy):
+  """A distribution strategy for synchronous training on multiple workers.
+
+  This strategy implements synchronous distributed training across multiple
+  workers, each with potentially multiple GPUs. Similar to
+  `tf.distribute.MirroredStrategy`, it replicates all variables and computations
+  to each local device. The difference is that it uses a distributed collective
+  implementation (e.g. all-reduce), so that multiple workers can work together.
+  """
+
+  def __init__(self, cluster_resolver=None, communication_options=None, *,
+               mesh=None):
+    """Creates the strategy.
+
+    Args:
+      cluster_resolver: optional
+        `tf.distribute.cluster_resolver.ClusterResolver`. In case neither `mesh`
+        nor `cluster_resolver` are provided,
+        `tf.distribute.cluster_resolver.TFConfigClusterResolver` is used.
+      communication_options: currently ignore.
+      mesh: optional Dtensor global mesh for the computation. Note that either
+        `mesh` or the `cluster_resolver` should be provided. and not both.
+    """
+    self._validate_init_args(mesh, cluster_resolver)
+    if not mesh:
+      if not cluster_resolver:
+        # Use the TFConfigClusterResolver as default
+        cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+      dtensor_env_var = _parse_dtensor_env_var_from_cluster_resolver(
+          cluster_resolver)
+      _config_dtensor_env_var(dtensor_env_var)
+      mesh = _build_distributed_mesh(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME)
+    extended = dtensor_strategy_extended.DTensorStrategyExtended(
+        container_strategy=self, mesh=mesh)
+    super().__init__(extended)
+    self._mesh = mesh
+    self._cluster_resolver = cluster_resolver
+
+  @classmethod
+  def _validate_init_args(cls, mesh, cluster_resolver):
+    if mesh and cluster_resolver:
+      raise ValueError('Mesh and cluster_resolver can not be provided at the '
+                       f'same time. Received mesh = {mesh}, cluster_resolver = '
+                       f'{cluster_resolver}')
+    if mesh and len(mesh.shape()) != 1:
+      raise ValueError('The mesh for MultiWorkerMirroredStrategy must be 1D, '
+                       f'received: {len(mesh.shape())}D')
+
+  def reduce(self, reduce_op, value, axis):
+    return dtensor_util.dtensor_reduce(self, reduce_op, value, axis)
+
+  @property
+  def mesh(self):
+    """Returns the mesh used by the strategy."""
+    return self._mesh
+
+
+def _parse_dtensor_env_var_from_cluster_resolver(cluster_resolver):
+  """Parse the env vars for Dtensor based on the cluster resolver.
+
+  In the multi-client setting, each of the DTensor jobs need to aware of each
+  other, and the interface to setup those values are via the envvars. The
+  value used by dtensor are different from the existing
+  `MultiWorkerMirroredStrategy`. This function will parse the value from
+  cluster resolver, and populate the corresponding value for DTensor jobs in the
+  `os.environ`.
+
+  Args:
+    cluster_resolver: A `tf.distribute.cluster_resolver.ClusterResolver`
+      instance.
+
+  Returns:
+    A dict of {Str:Str} which contains all the env vars needed by DTensor jobs.
+    The value is for verification purpose.
+
+  Raises:
+    The value parsed from existing cluster spec is not valid.
+  """
+  result = {}
+
+  # Retrieve the number of host, cluster config from the resolver.
+  cluster_spec = multi_worker_util.normalize_cluster_spec(
+      cluster_resolver.cluster_spec())
+  # Export all the necessary envvars for dtensor
+  # Get all the jobs from the cluster spec. Note that the in the normal
+  # setting, it could be multiple worker devices without chief, and the
+  # worker 0 will be the chief, or an explicit chief with multiple worker job.
+  dtensor_jobs = []
+  if 'chief' in cluster_spec.jobs:
+    dtensor_jobs.extend(cluster_spec.job_tasks('chief'))
+  if 'worker' in cluster_spec.jobs:
+    dtensor_jobs.extend(cluster_spec.job_tasks('worker'))
+
+  if None in dtensor_jobs:
+    raise ValueError('Unexpected dtensor job address from cluster spec: '
+                     f'{cluster_spec}')
+  result['DTENSOR_JOBS'] = ','.join(dtensor_jobs)
+  result['DTENSOR_NUM_CLIENTS'] = str(len(dtensor_jobs))
+
+  if cluster_resolver.task_type == 'chief':
+    dtensor_client_id = 0
+  elif cluster_resolver.task_type == 'worker':
+    dtensor_client_id = cluster_resolver.task_id
+    if 'chief' in cluster_spec.jobs:
+      dtensor_client_id += 1
+  result['DTENSOR_CLIENT_ID'] = str(dtensor_client_id)
+  result['DTENSOR_JOB_NAME'] = 'worker'
+
+  return result
+
+
+def _config_dtensor_env_var(dtensor_env_vars):
+  for k, v in dtensor_env_vars.items():
+    os.environ[k] = v
+
+
+def _build_distributed_mesh(batch_dim_name):
+  device_type = d_config.preferred_device_type()
+  local_devices = d_config.local_devices(device_type)
+  number_clients = d_config.num_clients()
+  dtensor_util.initialize_accelerator_system_once(device_type)
+  # This assumes each client has same number of devices.
+  mesh_dims = [(batch_dim_name, len(local_devices) * number_clients)]
+  return mesh_util.create_distributed_mesh(
+      mesh_dims, device_type=device_type)

tensorflow/python/distribute/experimental/multi_worker_mirrored_strategy_test.py

@@ -0,0 +1,647 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test for MultiWorkerMirroredStrategy backed by DTensor API."""
+
+import json
+import os
+
+from absl import flags
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import d_variable
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python.tests import multi_client_test_util
+from tensorflow.dtensor.python.tests import test_backend_util
+from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.distribute.experimental import multi_worker_mirrored_strategy as mwms
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test as tf_test
+
+
+class MultiWorkerMirroredStrategyTest(tf_test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.num_client = flags.FLAGS.num_clients
+    self.num_local_devices = flags.FLAGS.num_local_devices
+
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    self.client_id = int(tf_config['task']['index'])
+
+  def test_strategy_creation_with_default_cluster_resolver(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    mesh = strategy.mesh
+    self.assertIsNotNone(mesh)
+    self.assertLen(mesh.global_device_ids(),
+                   self.num_client * self.num_local_devices)
+    self.assertLen(mesh.local_device_ids(), self.num_local_devices)
+    self.assertIsInstance(strategy._cluster_resolver,
+                          tfconfig_cluster_resolver.TFConfigClusterResolver)
+
+  def test_invalid_init_arguments(self):
+    mesh = object()
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'Mesh and cluster_resolver can not be provided at the same time'):
+      mwms.MultiWorkerMirroredStrategy(
+          mesh=mesh,
+          cluster_resolver=cluster_resolver)
+
+  def test_parse_dtensor_env_var_from_cluster_resolver(self):
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+
+    dtensor_env_vars = mwms._parse_dtensor_env_var_from_cluster_resolver(
+        cluster_resolver)
+
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    worker_jobs = ','.join(tf_config['cluster']['worker'])
+    client_id = tf_config['task']['index']
+
+    self.assertLen(dtensor_env_vars, 4)
+    self.assertEqual(dtensor_env_vars['DTENSOR_JOBS'], worker_jobs)
+    self.assertEqual(dtensor_env_vars['DTENSOR_NUM_CLIENTS'],
+                     str(self.num_client))
+    self.assertEqual(dtensor_env_vars['DTENSOR_CLIENT_ID'], client_id)
+    self.assertEqual(dtensor_env_vars['DTENSOR_JOB_NAME'], 'worker')
+
+  @parameterized.named_parameters([
+      ('py_floats', lambda: [1.0, 2.0], True),
+      ('np_floats', lambda: np.array([1.0, 2.0]), True),
+      ('tf_const', lambda: constant_op.constant([1.0, 2.0]), True),
+      ('py_floats_callable', lambda: [1.0, 2.0], False),
+      ('np_floats_callable', lambda: np.array([1.0, 2.0]), False),
+      ('tf_const_callable', lambda: constant_op.constant([1.0, 2.0]), False),
+  ])
+  def test_variable_creation(self, init_value, convert_callable):
+    if convert_callable:
+      init_value = init_value()
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+      v = variables.Variable(init_value)
+
+    self.assertIsInstance(v, d_variable.DVariable)
+    self.assertIsNotNone(v.layout)
+    self.assertEqual(v.layout, layout.Layout.replicated(strategy.mesh, rank=1))
+
+  def test_strategy_extension(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertIsInstance(strategy.extended, distribute_lib.StrategyExtendedV2)
+
+  def test_num_replica_in_sync(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertEqual(strategy.num_replicas_in_sync,
+                     self.num_client * self.num_local_devices)
+
+  def test_mesh(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertIsNotNone(strategy.mesh)
+
+  def test_worker_devices(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    worker_devices = strategy.extended.worker_devices
+    self.assertLen(worker_devices, self.num_local_devices)
+    self.assertEqual(worker_devices, tuple(strategy.mesh.local_devices()))
+
+  def test_parameter_devices(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    parameter_devices = strategy.extended.parameter_devices
+    self.assertLen(parameter_devices, self.num_local_devices)
+    self.assertEqual(parameter_devices, tuple(strategy.mesh.local_devices()))
+
+  def test_variable_created_in_scope(self):
+    strategy1 = mwms.MultiWorkerMirroredStrategy()
+    with strategy1.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    v2 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    strategy2 = mwms.MultiWorkerMirroredStrategy()
+    with strategy2.scope():
+      v3 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    self.assertTrue(strategy1.extended.variable_created_in_scope(v1))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v2))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v3))
+    self.assertTrue(strategy2.extended.variable_created_in_scope(v3))
+
+  def test_colocate_vars_with(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+      with strategy.extended.colocate_vars_with(v1):
+        v2 = variables.Variable(constant_op.constant([2.0, 3.0]))
+
+    # We assert the layout for the variable, and make sure they are same.
+    self.assertEqual(v1.layout, v2.layout)
+
+  def test_in_multi_worker_mode(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertTrue(strategy.extended._in_multi_worker_mode())
+
+  def test_run_with_distribute_value_input(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    def value_fn(value_context):
+      return value_context.replica_id_in_sync_group
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, self.num_local_devices)
+    # Note that the scalar value from
+    # experimental_distribute_values_from_function will be up rank to 1D since
+    # batched shared dtensor need at least be 1D.
+
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result.values[i],
+          constant_op.constant(
+              [(self.client_id * self.num_local_devices + i) * 2]))
+
+  def test_nested_structure_output(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    def value_fn(ctx):
+      value = float(ctx.num_replicas_in_sync)
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertLen(result.keys(), 2)
+    self.assertIsInstance(result['a'], dtensor_util.DTensorDistributedValue)
+    self.assertLen(result['a'].values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result['a'].values[i],
+          constant_op.constant([strategy.num_replicas_in_sync * 2.0]))
+
+    self.assertIsInstance(result['b'], dtensor_util.DTensorDistributedValue)
+    self.assertLen(result['b'].values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result['b'].values[i],
+          constant_op.constant([(strategy.num_replicas_in_sync + 1.0) * 2.0,
+                                (strategy.num_replicas_in_sync + 2.0) * 2.0]))
+
+  def test_inputs_with_dtensor_distribute_values(self):
+
+    @def_function.function
+    def replica_fn_1(inputs):
+      return inputs * 2.0
+
+    @def_function.function
+    def replica_fn_2(inputs):
+      return inputs + 1.0
+
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant(3.0)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    result_1 = strategy.run(replica_fn_1, args=(d_tensor_input,))
+    self.assertIsInstance(result_1, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_1.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(result_1.values[i], constant_op.constant([6.0]))
+
+    result_2 = strategy.run(replica_fn_2, args=(result_1,))
+    self.assertIsInstance(result_2, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_2.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(result_2.values[i], constant_op.constant([7.0]))
+
+  def test_get_replica_context(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    tensor_input = constant_op.constant(3)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    @def_function.function
+    def replica_fn(inputs):
+      replica_context = distribute_lib.get_replica_context()
+      self.assertIsInstance(replica_context, dtensor_util.DTensorReplicaContext)
+      return inputs * replica_context.num_replicas_in_sync
+
+    # Default replica context
+    self.assertIsNotNone(distribute_lib.get_replica_context())
+    with strategy.scope():
+      self.assertIsNone(distribute_lib.get_replica_context())
+
+      result = strategy.run(replica_fn, args=(d_tensor_input,))
+
+    self.assertLen(result.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result.values[i],
+          constant_op.constant([3 * strategy.num_replicas_in_sync]))
+
+  def test_gather_non_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant(3.0)
+
+    result = strategy.gather(tensor_input, axis=0)
+    self.assertAllClose(result, tensor_input)
+
+  def test_gather_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    stride = self.num_client * self.num_local_devices
+
+    def value_fn(value_context):
+      start = value_context.replica_id_in_sync_group * stride
+      return array_ops.reshape(
+          math_ops.range(start=start, limit=start + stride), shape=(1, stride)
+      )
+    distribute_result = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+    # distribute_result is a DTensorDistributedValue.
+    # The shape of the global tensor is [stride, stride],
+    # and each worker gets [stride/2, stride].
+    result = strategy.gather(distribute_result, axis=0)
+
+    start = stride * self.num_local_devices * self.client_id
+    end = start + stride * self.num_local_devices
+    self.assertEqual(result.shape, [self.num_local_devices, stride])
+    self.assertAllClose(
+        result,
+        array_ops.reshape(
+            math_ops.range(start=start, limit=end),
+            shape=(self.num_local_devices, -1),
+        ),
+    )
+
+    result = strategy.gather(distribute_result, axis=1)
+    self.assertEqual(result.shape, [1, self.num_local_devices * stride])
+    self.assertAllClose(
+        result,
+        array_ops.reshape(
+            math_ops.range(start=start, limit=end), shape=(1, -1)
+        ),
+    )
+
+  def test_reduce_mean_non_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'
+    ):
+      strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=0)
+
+  def test_reduce_sum_non_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'
+    ):
+      strategy.reduce(reduce_util.ReduceOp.SUM, tensor_input, axis=0)
+
+  def test_reduce_mean_distribute_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * 4.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+    # replica 0 has [[0.0, 1.0],[2.0, 3.0]] and
+    # replica 1 has [[4.0, 5.0],[6.0, 7.0]]. Each worker has 4 replicas.
+    # For worker 2, it has replica 4 ~ 7.
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=None
+    )
+    # This should be a global reduce and each worker should have same value.
+    # [[14.0, 15.0],[16.0, 17.0]]
+    final = (self.num_local_devices * self.num_client - 1) * 2.0
+    self.assertAllClose(
+        result, constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + final
+    )
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=0
+    )
+    # [15.0, 16.0]
+    self.assertAllClose(result, constant_op.constant([0.0, 1.0]) + final + 1)
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=1
+    )
+
+    self.assertAllClose(result, constant_op.constant([0.5, 2.5]) + final)
+
+  def test_reduce_sum_distribute_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * 4.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+    # replica 0 has [[0.0, 1.0],[2.0, 3.0]] and
+    # replica 1 has [[4.0, 5.0],[6.0, 7.0]]. Each worker has 4 replicas.
+    # For worker 2, it has replica 4 ~ 7.
+    # The shape of the global tensor is [16, 2], and each worker gets [8, 2].
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=None
+    )
+    self.assertAllClose(result, [[112.0, 120.0], [128.0, 136.0]])
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, distribute_value, axis=0)
+    self.assertAllClose(result, constant_op.constant([240.0, 256.0]))
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, distribute_value, axis=1)
+    self.assertAllClose(result, constant_op.constant([232.0, 264.0]))
+
+  def test_reduce_mean_mirrored_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([2.0, 3.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([1.5, 3.5]))
+
+  def test_reduce_sum_mirrored_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([4.0, 6.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([3.0, 7.0]))
+
+  def test_reduce_value_device(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=None)
+    self.assertIn('CPU:0', result.device)
+
+  def test_experimental_local_results(self):
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([3.0])
+
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    result = strategy.run(replica_fn)
+    local_result = strategy.experimental_local_results(result)
+
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertEqual(local_result[i], constant_op.constant([3.0]))
+
+  def test_experimental_local_results_with_inputs(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    def value_fn(ctx):
+      value = float(ctx.num_replicas_in_sync)
+      return {'a': value, 'b': constant_op.constant([value + 1.0, value + 2.0])}
+
+    distributed_values = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    local_result = strategy.experimental_local_results(result)
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertDictEqual(
+          local_result[i],
+          {
+              'a': constant_op.constant([16.0]),
+              'b': constant_op.constant([18.0, 20.0]),
+          },
+      )
+
+
+class StrategyDatasetTest(tf_test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.num_client = flags.FLAGS.num_clients
+    self.num_local_devices = flags.FLAGS.num_local_devices
+
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    self.client_id = int(tf_config['task']['index'])
+
+    self.images = stateless_random_ops.stateless_random_uniform(
+        [8, 8, 3], seed=(1, 2), minval=0, maxval=255)
+    self.labels = stateless_random_ops.stateless_random_uniform(
+        [1], seed=(1, 2), minval=0, maxval=10)
+
+    self.dataset = dataset_ops.Dataset.from_tensors(
+        (self.images, self.labels)).repeat()
+
+  def test_create_batched_dataset(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    global_batch_size = self.num_client * self.num_local_devices * 2
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    element = next(iter(distributed_dataset))
+    batched_image, batched_label = element
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # After unpack, it should only get the local shards.
+    self.assertLen(d_api.unpack(batched_image), self.num_local_devices)
+    self.assertLen(d_api.unpack(batched_label), self.num_local_devices)
+
+  def test_uneven_batched_dataset(self):
+    elements = [[1, 2, 3], [1, 2], [1, 2, 3, 4]]
+    dataset = dataset_ops.Dataset.from_generator(
+        lambda: elements, dtypes.int64).repeat()
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with self.assertRaisesRegex(ValueError, 'requires a static batch size'):
+      strategy.experimental_distribute_dataset(dataset)
+
+  def test_deprecated_strategy_methods(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_dataset_iterator(self.dataset)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_input_fn_iterator(lambda _: self.dataset)
+
+  def test_distribute_dataset_from_fn(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    local_batch_size = 4
+    global_batch_size = local_batch_size * strategy.num_replicas_in_sync
+
+    def dataset_fn(option):
+      del option
+      return dataset_ops.Dataset.from_tensors(
+          (self.images, self.labels)).repeat().batch(
+              local_batch_size, drop_remainder=True).prefetch(2)
+
+    distributed_dataset = strategy.distribute_datasets_from_function(
+        dataset_fn, None)
+    iterator = iter(distributed_dataset)
+
+    self.assertEqual(distributed_dataset.element_spec,
+                     (tensor_spec.TensorSpec(shape=(global_batch_size, 8, 8, 3),
+                                             dtype=dtypes.float32, name=None),
+                      tensor_spec.TensorSpec(shape=(global_batch_size, 1),
+                                             dtype=dtypes.float32, name=None)))
+    self.assertEqual(distributed_dataset.element_spec, iterator.element_spec)
+
+    batched_image, batched_label = next(iterator)
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # After unpack, it should only get the local shards.
+    unpacked_images = d_api.unpack(batched_image)
+    self.assertLen(unpacked_images, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertEqual(unpacked_images[i].shape, [local_batch_size, 8, 8, 3])
+
+  def test_distribute_values_from_function(self):
+    array_value = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
+    def value_fn(ctx):
+      return array_value[ctx.replica_id_in_sync_group]
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    self.assertEqual(d_api.fetch_layout(distributed_values),
+                     layout.Layout.batch_sharded(
+                         strategy.mesh, batch_dim='batch', rank=1))
+    unpacked_value = d_api.unpack(distributed_values)
+    self.assertLen(unpacked_value, self.num_local_devices)
+    start = 1.0 + self.num_local_devices * self.client_id
+    for i in range(self.num_local_devices):
+      self.assertEqual(unpacked_value[i], start + i)
+
+  def test_distribute_dataset_in_tf_function(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    local_batch_size = 4
+    global_batch_size = local_batch_size * strategy.num_replicas_in_sync
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+
+    @def_function.function
+    def step_fn(iterator):
+      images, labels = next(iterator)
+      del labels
+      return images
+
+    result = strategy.run(step_fn, args=(iter(distributed_dataset),))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertEqual(result.values[i].shape, [local_batch_size, 8, 8, 3])
+
+
+def client_config_function(config_params):
+  client_id = config_params['client_id']
+  worker_jobs = config_params['worker_jobs']
+  num_devices = config_params['num_devices']
+
+  os.environ['TF_CONFIG'] = json.dumps({
+      'cluster': {
+          'worker': worker_jobs
+      },
+      'task': {'type': 'worker', 'index': f'{client_id}'}
+  })
+
+  if config.list_physical_devices('GPU'):
+    device_type = 'GPU'
+  elif test_util.is_tpu_present():
+    device_type = 'TPU'
+  else:
+    device_type = 'CPU'
+
+  # reset_logical_devices
+  test_util.reset_context()
+  if device_type != 'TPU':
+    # Configure virtual devices. This does not initialize the TensorFlow
+    # context.
+    test_util.reset_logical_devices(device_type, num_devices)
+
+  # Validates the correct number of devices are created.
+  logical_devices = test_util.list_local_logical_devices(device_type)
+  assert len(logical_devices) == num_devices, (
+      logical_devices,
+      f'Test is mis-configured: expecting {num_devices} logical_devices.')
+
+
+if __name__ == '__main__':
+  test_backend_util.handle_test_main(
+      multi_client_test_util.multi_client_main, client_config_function)

tensorflow/python/distribute/experimental/rpc/BUILD

@@ -0,0 +1,64 @@
+# Python bindings for RPC client and server ops.
+
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "rpc_ops",
+    srcs = [
+        "rpc_ops.py",
+    ],
+    deps = [
+        "//tensorflow/distribute/experimental/rpc/kernels:gen_rpc_ops",
+        "//tensorflow/distribute/experimental/rpc/proto:tf_rpc_service_proto_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:none_tensor",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "rpc_ops_test",
+    size = "medium",
+    srcs = ["rpc_ops_test.py"],
+    shard_count = 7,
+    tags = [
+        "no_mac",  # flaky, b/205156709
+    ],
+    deps = [
+        ":rpc_ops",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:data_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+        "@pypi//portpicker",
+    ],
+)

tensorflow/python/distribute/experimental/rpc/rpc_ops.py

@@ -0,0 +1,505 @@
+# 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.
+# ==============================================================================
+"""Module to expose RPC APIs in tensorflow."""
+
+from typing import Optional, Sequence, Union
+
+import tensorflow.distribute.experimental.rpc.kernels.gen_rpc_ops as gen_rpc_ops
+from tensorflow.distribute.experimental.rpc.proto import tf_rpc_service_pb2 as rpc_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as tf_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import none_tensor
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.types import core as core_tf_types
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+def get_output_specs_from_function(func: tf_function.ConcreteFunction):
+  output_specs = nest.map_structure(type_spec.type_spec_from_value,
+                                    func.structured_outputs)
+  output_specs_proto = nested_structure_coder.encode_structure(output_specs)
+  return output_specs_proto.SerializeToString()
+
+
+def get_input_specs_from_function(func: tf_function.ConcreteFunction):
+  arg_specs, _ = func.structured_input_signature
+  arg_specs_proto = nested_structure_coder.encode_structure(arg_specs)
+  return arg_specs_proto.SerializeToString()
+
+
+@tf_export("distribute.experimental.rpc.Server", v1=[])
+class Server(object):
+  """A Server base class for accepting RPCs for registered tf.functions.
+
+    Functions can be registered on the server and are exposed via RPCs.
+  """
+
+  @staticmethod
+  def create(rpc_layer, address):
+    """Create TF RPC server at given address.
+
+    Args:
+      rpc_layer: Communication layer between client and server. Only "grpc" rpc
+        layer is supported at the moment.
+      address: Address where RPC server is hosted.
+
+    Returns:
+      An instance of `tf.distribute.experimental.rpc.Server` class.
+
+    Raises:
+        A ValueError if rpc_layer other than "grpc" is used. Only GRPC
+        is supported at the moment.
+
+    Example usage:
+
+      >>> import portpicker
+      >>> @tf.function(input_signature=[
+      ...      tf.TensorSpec([], tf.int32),
+      ...      tf.TensorSpec([], tf.int32)])
+      ... def remote_fn(a, b):
+      ...   return tf.add(a, b)
+
+      >>> port = portpicker.pick_unused_port()
+      >>> address = "localhost:{}".format(port)
+      >>> server = tf.distribute.experimental.rpc.Server.create("grpc", address)
+      >>> server.register("addition", remote_fn)
+      >>> server.start()
+
+    """
+    if rpc_layer != "grpc":
+      raise ValueError("Only GRPC backend is supported at the moment.")
+    return GrpcServer(address=address)
+
+  def register(self, method_name: str,
+               func: Union[def_function.Function,
+                           tf_function.ConcreteFunction]):
+    """Method for registering tf.function on server.
+
+    Registered methods can be invoked remotely from clients.
+
+    Args:
+      method_name: Name of the tf.function. Clients use this method_name to make
+        RPCs.
+      func: A `tf.function` or ConcreteFunction to register.
+    """
+    raise NotImplementedError("Please use create_server method to create a"
+                              "concrete subclass of Server.")
+
+  def start(self):
+    """Starts the RPC server on provided address.
+
+     Server listens for new requests from client, once it is started.
+    """
+    raise NotImplementedError("Please use create_server method to create a"
+                              "concrete subclass of Server.")
+
+
+@tf_export("distribute.experimental.rpc.Client", v1=[])
+class Client(object):
+  """Client class for invoking RPCs to the server."""
+
+  @staticmethod
+  def create(rpc_layer, address, name="", timeout_in_ms=0):
+    """Create TF RPC client to connect to the given address.
+
+    Args:
+      rpc_layer: Communication layer between client and server. Only "grpc" rpc
+        layer is supported at the moment.
+      address: Address of the server to connect the RPC client to.
+      name: Name of the RPC Client. You can create multiple clients connecting
+        to same server and distinguish them using different names.
+      timeout_in_ms: The default timeout to use for outgoing RPCs from client. 0
+        indicates no timeout. Exceeding timeout during RPC will raise
+        DeadlineExceeded error.
+
+    Returns:
+      An instance of `tf.distribute.experimental.rpc.Client` with the following
+      dynamically added methods for eagerly created clients:
+        * `Registered methods` e.g. multiply(**args):
+            If Client is created when executing eagerly, client will request the
+            list of registered methods from server during client creation.
+            The convenience methods for RPCs will be dynamically added to the
+            created Client instance.
+
+            For example, when a server has method "multiply" registered, the
+            client object created in eager mode will have 'multiply' method
+            available. Users can use client.multiply(..) to make RPC, instead of
+            client.call("multiply", ...)
+
+            Both "call" and "multiply" methods are non-blocking i.e. they return
+            a StatusOrResult object which should be used to wait for getting
+            value or error.
+
+            Along with the above, blocking versions of the registered
+            methods are also dynamically added to client instance.
+            e.g. multiply_blocking(**args). These methods block till the RPC is
+            finished and return response for successful RPC. Otherwise raise
+            exception.
+
+            These methods are not available when Client is created inside a
+            tf.function.
+
+    Raises:
+        A ValueError if rpc_layer other than "grpc" is used. Only GRPC
+          is supported at the moment.
+        A DeadlineExceeded exception in eager mode if timeout exceeds while
+          creating and listing client methods.
+
+    Example usage:
+      >>> # Have server already started.
+      >>> import portpicker
+      >>> @tf.function(input_signature=[
+      ...      tf.TensorSpec([], tf.int32),
+      ...      tf.TensorSpec([], tf.int32)])
+      ... def remote_fn(a, b):
+      ...   return tf.add(a, b)
+
+      >>> port = portpicker.pick_unused_port()
+      >>> address = "localhost:{}".format(port)
+      >>> server = tf.distribute.experimental.rpc.Server.create("grpc", address)
+      >>> server.register("addition", remote_fn)
+      >>> server.start()
+
+      >>> # Start client
+      >>> client = tf.distribute.experimental.rpc.Client.create("grpc",
+      ...      address=address, name="test_client")
+
+      >>> a = tf.constant(2, dtype=tf.int32)
+      >>> b = tf.constant(3, dtype=tf.int32)
+
+      >>> result = client.call(
+      ...    args=[a, b],
+      ...    method_name="addition",
+      ...    output_specs=tf.TensorSpec((), tf.int32))
+
+      >>> if result.is_ok():
+      ...   result.get_value()
+
+      >>> result = client.addition(a, b)
+
+      >>> if result.is_ok():
+      ...   result.get_value()
+
+      >>> value = client.addition_blocking(a, b)
+    """
+    if rpc_layer != "grpc":
+      raise ValueError("Only GRPC backend is supported at the moment.")
+    if context.executing_eagerly():
+      list_registered_methods = True
+    else:
+      list_registered_methods = False
+    return GrpcClient(
+        address=address,
+        name=name,
+        list_registered_methods=list_registered_methods,
+        timeout_in_ms=timeout_in_ms)
+
+  def call(self,
+           method_name: str,
+           args: Optional[Sequence[core_tf_types.Tensor]] = None,
+           output_specs=None,
+           timeout_in_ms=0):
+    """Method for making RPC calls to remote server.
+
+    This invokes RPC to the server, executing the registered method_name
+    remotely.
+    Args:
+      method_name: Remote registered method to invoke
+      args: List of arguments for the registered method.
+      output_specs: Output specs for the output from method.
+         For example, if tf.function is: @tf.function(input_signature=[
+           tf.TensorSpec([], tf.int32), tf.TensorSpec([], tf.int32) ])
+          def multiply_fn(a, b): return tf.math.multiply(a, b)
+        output_spec is: tf.TensorSpec((), tf.int32)  If you have access to TF
+          Function, the output specs can be generated
+       from tf.function by calling: output_specs =
+         tf.nest.map_structure(tf.type_spec_from_value,
+         tf_function.get_concrete_function().structured_outputs  If output_specs
+         are not provided, flattened list of tensors will be returned in
+         response.
+      timeout_in_ms: Timeout for this call. If 0, default client timeout will be
+        used.
+
+    Returns:
+      An instance of `StatusOrResult` class with the following available
+      methods.
+        * `is_ok()`:
+            Returns True of RPC was successful.
+        * `get_error()`:
+            Returns TF error_code and error message for the RPC.
+        * `get_value()`:
+            Returns the returned value from remote TF function execution
+            when RPC is successful.
+
+      Calling any of the above methods will block till RPC is completed and
+      result is available.
+    """
+    raise NotImplementedError("Must be implemented in inherited classes.")
+
+
+class GrpcServer(Server):
+  """GrpcServer object encapsulates a resource with GRPC server.
+
+    Functions can be registered locally and are exposed via RPCs.
+    Example:
+    ```
+    server = rpc_ops.GrpcServer("host:port")
+    @tf.function
+    def add(a, b):
+      return a + b
+
+    server.register("add", add)
+    server.start()
+    ```
+  """
+
+  def __init__(self, address: str):
+    self._server_handle = gen_rpc_ops.rpc_server(address)
+    if context.executing_eagerly():
+      self._handle_deleter = resource_variable_ops.EagerResourceDeleter(
+          handle=self._server_handle, handle_device=self._server_handle.device)
+    else:
+      raise NotImplementedError("Please create the server outside tf.function.")
+
+  def register(self, method_name: str,
+               func: Union[def_function.Function,
+                           tf_function.ConcreteFunction]):
+    """Method for registering functions."""
+
+    if isinstance(func, def_function.Function):
+      if func.function_spec.arg_names:
+        if func.input_signature is None:
+          raise ValueError("Input signature not specified for the function.")
+      concrete_fn = func.get_concrete_function()
+      gen_rpc_ops.rpc_server_register(
+          self._server_handle,
+          method_name=method_name,
+          captured_inputs=concrete_fn.captured_inputs,
+          input_specs=get_input_specs_from_function(concrete_fn),
+          output_specs=get_output_specs_from_function(concrete_fn),
+          f=concrete_fn)
+    elif isinstance(func, tf_function.ConcreteFunction):
+      gen_rpc_ops.rpc_server_register(
+          self._server_handle,
+          method_name=method_name,
+          captured_inputs=func.captured_inputs,
+          input_specs=get_input_specs_from_function(func),
+          output_specs=get_output_specs_from_function(func),
+          f=func)
+    else:
+      # Python functions
+      # TODO(b/186762191): Add an implementation to support python functions.
+      raise ValueError("Only TF functions are supported with Register method")
+
+  def start(self):
+    """Starts GRPC server."""
+    gen_rpc_ops.rpc_server_start(self._server_handle)
+
+
+class GrpcClient(Client):
+  """Client wrapper to connect to remote RPC server using GRPC.
+
+  If Client is created with (list_registered_methods=True):
+  1. Input and output specs for the methods till this point will be fetched from
+  Server.
+  2. convenience methods are added to invoke registered methods directly from
+  client.
+  For example:
+    For call a server method `add`
+    client.add(a, b) or client.add_async(a, b) can be used instead of
+    client.call(args=[a,b], output_specs=[..])
+
+  Prerequisite for using list_registered_methods=True:
+   1. Server should be already started with the registered methods.
+   2. Client must be created in Eager mode.
+  """
+
+  def __init__(self,
+               address: str,
+               name: str = "",
+               list_registered_methods=False,
+               timeout_in_ms=0):
+    self._client_handle, methods = gen_rpc_ops.rpc_client(
+        shared_name=name,
+        server_address=address,
+        list_registered_methods=list_registered_methods,
+        timeout_in_ms=timeout_in_ms)
+    if context.executing_eagerly():
+      self._handle_deleter = resource_variable_ops.EagerResourceDeleter(
+          handle=self._client_handle, handle_device=self._client_handle.device)
+    else:
+      raise NotImplementedError(
+          "Client creation is supported only in eager mode.")
+    self._server_address = address
+    self._method_registry = {}
+    for method in methods.numpy():
+      m = rpc_pb2.RegisteredMethod()
+      m.ParseFromString(method)
+      output_specs = nested_structure_coder.decode_proto(m.output_specs)
+      input_specs = nested_structure_coder.decode_proto(m.input_specs)
+      self._method_registry[m.method] = output_specs
+      # TODO(ishark): Perhaps doc string can also be taken as input during
+      # function registration.
+      doc_string = "RPC Call for " + m.method + " method to server " + address
+      self._add_method(m.method, output_specs, input_specs, self._client_handle,
+                       doc_string)
+
+  def _add_method(self, method_name, output_specs, input_specs, client_handle,
+                  doc_string):
+    """Method to add RPC methods to the client object."""
+
+    def validate_and_get_flat_inputs(*args):
+      if args is None:
+        args = []
+      if input_specs:
+        nest.assert_same_structure(args, input_specs)
+      flat_inputs = nest.flatten(args)
+      return flat_inputs
+
+    def call_wrapper(*args, timeout_in_ms=0):
+      status_or, deleter = gen_rpc_ops.rpc_call(
+          client_handle,
+          args=validate_and_get_flat_inputs(*args),
+          method_name=method_name,
+          timeout_in_ms=timeout_in_ms)
+      return StatusOrResult(status_or, deleter, output_specs)
+
+    def call_blocking_wrapper(*args, timeout_in_ms=0):
+      status_or, deleter = gen_rpc_ops.rpc_call(
+          client_handle,
+          args=validate_and_get_flat_inputs(*args),
+          method_name=method_name,
+          timeout_in_ms=timeout_in_ms)
+      status_or = StatusOrResult(status_or, deleter, output_specs)
+      if status_or.is_ok():
+        return status_or.get_value()
+      else:
+        error_code, error_msg = status_or.get_error()
+        raise errors.exception_type_from_error_code(error_code.numpy())(
+            None, None, error_msg.numpy())
+
+    setattr(self, method_name, call_wrapper)
+    call_wrapper.__doc__ = doc_string
+
+    blocking_method_name = method_name + "_blocking"
+    setattr(self, blocking_method_name, call_blocking_wrapper)
+    call_blocking_wrapper.__doc__ = doc_string
+
+  def call(self,
+           method_name: str,
+           args: Optional[Sequence[core_tf_types.Tensor]] = None,
+           output_specs=None,
+           timeout_in_ms=0):
+    """Method to invoke remote registered functions on the connected server.
+
+    Server should be started before making an RPC Call.
+
+    Args:
+      method_name: Registered method to invoke on Server.
+      args: Input arguments for the method.
+      output_specs: Output specs for the output from method.
+      timeout_in_ms: Timeout for this call. If 0, default client timeout will be
+       used.
+
+    Returns:
+      StatusOrResult object. This function issues the RPC call to server, it
+      does not block for the duration of RPC. Please call is_ok, get_error or
+      get_value methods on the returned object to blocked till RPC finishes.
+    """
+    if args is None:
+      args = []
+    status_or, deleter = gen_rpc_ops.rpc_call(
+        self._client_handle,
+        args=nest.flatten(args),
+        method_name=method_name,
+        timeout_in_ms=timeout_in_ms)
+    return StatusOrResult(status_or, deleter, output_specs)
+
+
+class StatusOrResult(object):
+  """Class representing result and status from RPC Call."""
+
+  def __init__(self, status_or, deleter, output_specs=None):
+    self._status_or = status_or
+    self._output_specs = output_specs
+    self._deleter = deleter
+    self._error_code: dtypes.int64 = None
+    self._error_message: dtypes.string = None
+
+  def _check_status(self):
+    if self._error_code is None:
+      self._error_code, self._error_message = gen_rpc_ops.rpc_check_status(
+          self._status_or)
+
+  def __del__(self):
+    # Make sure the resource is deleted in the same mode as it was created in.
+    if context.executing_eagerly():
+      with context.eager_mode():
+        gen_rpc_ops.delete_rpc_future_resource(
+            handle=self._status_or, deleter=self._deleter)
+    else:
+      with context.graph_mode():
+        gen_rpc_ops.delete_rpc_future_resource(
+            handle=self._status_or, deleter=self._deleter)
+
+  def is_ok(self):
+    """Returns True if RPC is successful, otherwise returns False.
+
+    This call will block for RPC result.
+    """
+    self._check_status()
+    return math_ops.equal(self._error_code,
+                          constant_op.constant(0, dtype=dtypes.int64))
+
+  def get_error(self):
+    """Returns (TF Error Code, Error Message) from RPC Response.
+
+    This call will block for RPC result.
+    """
+    self._check_status()
+    return self._error_code, self._error_message
+
+  def get_value(self):
+    """Returns the returned response value from RPC Call when RPC is successful.
+
+      The returned value is tensors in the output_specs format as returned from
+      the RPC call
+
+
+    This call will block for RPC result.
+    """
+
+    self._check_status()
+    if self._output_specs is None or isinstance(self._output_specs,
+                                                none_tensor.NoneTensorSpec):
+      flat_output_dtypes = []
+      return_none = True
+    else:
+      return_none = False
+      flat_output_dtypes = [s.dtype for s in nest.flatten(self._output_specs)]
+
+    result = gen_rpc_ops.rpc_get_value(self._status_or, Tout=flat_output_dtypes)
+    if return_none:
+      return None
+    else:
+      return nest.pack_sequence_as(self._output_specs, result)

tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py

@@ -0,0 +1,850 @@
+# 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 rpc_ops.py."""
+
+import threading
+import time
+import numpy as np
+import portpicker
+
+from tensorflow.python.distribute.experimental.rpc import rpc_ops
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function as eager_def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import data_flow_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.platform import test
+from tensorflow.python.util import nest
+
+
+@test_util.with_eager_op_as_function
+class RpcOpsTest(test.TestCase):
+
+  def setUp(self):
+    super(RpcOpsTest, self).setUp()
+    cpus = config.list_physical_devices("CPU")
+    # Set 2 virtual CPUs
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+
+  def test_generated_rpc_ops(self):
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    concrete_remote_fn = remote_fn.get_concrete_function()
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.gen_rpc_ops.rpc_server(server_address=address)
+
+    rpc_ops.gen_rpc_ops.rpc_server_register(
+        server_resource,
+        f=concrete_remote_fn,
+        captured_inputs=concrete_remote_fn.captured_inputs,
+        output_specs=rpc_ops.get_output_specs_from_function(concrete_remote_fn),
+        method_name="multiply")
+
+    rpc_ops.gen_rpc_ops.rpc_server_start(server_resource)
+    client_handle, _ = rpc_ops.gen_rpc_ops.rpc_client(
+        server_address=address, timeout_in_ms=5000)
+    future_resource, deleter = rpc_ops.gen_rpc_ops.rpc_call(
+        client_handle, args=[a, b], method_name="multiply", timeout_in_ms=0)
+
+    error_code, _ = rpc_ops.gen_rpc_ops.rpc_check_status(future_resource)
+    self.assertAllEqual(error_code, 0)
+    self.assertAllEqual(
+        rpc_ops.gen_rpc_ops.rpc_get_value(future_resource, Tout=[dtypes.int32]),
+        [6])
+
+    resource_variable_ops.EagerResourceDeleter(
+        handle=server_resource, handle_device=server_resource.device)
+
+    resource_variable_ops.EagerResourceDeleter(
+        handle=client_handle, handle_device=client_handle.device)
+
+    rpc_ops.gen_rpc_ops.delete_rpc_future_resource(future_resource, deleter)
+
+  def test_exported_rpc_api_static_factory(self):
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.Server.create("grpc", address)
+    server_resource.register("multiply", _remote_fn)
+
+    server_resource.start()
+    client = rpc_ops.Client.create("grpc", address=address, name="test_client")
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    mul_or = client.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    # Test empty client name
+    client1 = rpc_ops.Client.create("grpc", address)
+    mul_or = client1.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    # Test without output_spec
+    mul_or = client1.multiply(a, b)
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    self.assertEqual(client1.multiply.__doc__,
+                     "RPC Call for multiply method to server " + address)
+
+  def test_rpc_ops_call_method(self):
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def add_fn(a, b):
+      return math_ops.add(a, b)
+
+    # Register TF function
+    server_resource.register("multiply", _remote_fn)
+
+    # Register concrete Function
+    server_resource.register("add", add_fn.get_concrete_function())
+
+    server_resource.start()
+    client = rpc_ops.GrpcClient(address=address, name="test_client")
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    mul_or = client.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    add_or = client.call(
+        args=[a, b],
+        method_name="add",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+
+    self.assertAllEqual(add_or.is_ok(), True)
+    self.assertAllEqual(add_or.get_value(), 5)
+
+    # Test empty client name
+    client1 = rpc_ops.GrpcClient(address, list_registered_methods=True)
+    mul_or = client1.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+  def test_rpc_ops_non_blocking_convenience_methods(self):
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    # Register TF function
+    server_resource.register("multiply", _remote_fn)
+
+    server_resource.start()
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+
+    mul_or = client.multiply(a, b)
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    self.assertEqual(client.multiply.__doc__,
+                     "RPC Call for multiply method to server " + address)
+
+  def test_rpc_ops_blocking_convenience_methods(self):
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    # Register TF function
+    server_resource.register("multiply", _remote_fn)
+
+    server_resource.start()
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+    self.assertAllEqual(client.multiply_blocking(a, b), 6)
+
+    self.assertEqual(
+        client.multiply_blocking.__doc__,
+        "RPC Call for multiply method to server " + address)
+
+  def test_output_specs(self):
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def test_dict(val):
+      return {"key": val}
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def is_positive(a):
+      if a > 0:
+        return True
+      return False
+
+    @eager_def_function.function(input_signature=[])
+    def do_nothing():
+      return []
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def test_nested_structure(v):
+      return {"test": (v, [v, v]), "test1": (v,)}
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    server_resource.register("test_dict", test_dict)
+    server_resource.register("is_positive", is_positive)
+    server_resource.register("test_nested_structure", test_nested_structure)
+    server_resource.register("do_nothing", do_nothing)
+
+    server_resource.start()
+
+    client = rpc_ops.GrpcClient(
+        address=address, name="test_client", list_registered_methods=True)
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+
+    result_or = client.test_dict(a)
+    self.assertAllEqual(result_or.is_ok(), True)
+    nest.map_structure(self.assertAllEqual, result_or.get_value(), {"key": 2})
+
+    result_or = client.is_positive(a)
+    self.assertTrue(result_or.is_ok())
+    self.assertTrue(result_or.get_value())
+
+    result_or = client.test_nested_structure(a)
+    self.assertAllEqual(result_or.is_ok(), True)
+    nest.map_structure(self.assertAllEqual, result_or.get_value(), {
+        "test": (2, [2, 2]),
+        "test1": (2,)
+    })
+
+    result_or = client.do_nothing()
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [])
+
+  def test_input_specs(self):
+
+    @eager_def_function.function(input_signature=[{
+        "a": tensor_spec.TensorSpec([], dtypes.int32),
+        "b": tensor_spec.TensorSpec([], dtypes.int32)
+    }])
+    def test_input_dict(value):
+      return math_ops.add(value["a"], value["b"])
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    server_resource.register("test_input_dict", test_input_dict)
+
+    server_resource.start()
+
+    client = rpc_ops.GrpcClient(
+        address=address, name="test_client", list_registered_methods=True)
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+    result_or = client.test_input_dict({"a": a, "b": b})
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), 5)
+
+    with self.assertRaises(TypeError):
+      client.test_input_dict([a, b])
+
+  def test_call_register_ordering(self):
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+
+    # Create client succeeds before server start and registration
+    client = rpc_ops.GrpcClient(address)
+
+    # Create client with list_registered_methods fails before server is started.
+    with self.assertRaises(errors.DeadlineExceededError):
+      rpc_ops.GrpcClient(
+          address,
+          name="client1",
+          list_registered_methods=True,
+          timeout_in_ms=1)
+
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    server = rpc_ops.GrpcServer(address)
+
+    def start_server():
+      # Delay server start to test whether client creation also waits
+      # till server is up.
+      time.sleep(1)
+      server.register("assign_add", assign_add)
+      server.start()
+
+    t = threading.Thread(target=start_server)
+    t.start()
+
+    # Create same "client1" again should succeed.
+    client1_with_listed_methods = rpc_ops.GrpcClient(
+        address, name="client1", list_registered_methods=True)
+
+    result_or = client1_with_listed_methods.assign_add(
+        variables.Variable(2, dtype=dtypes.int64))
+    self.assertAllEqual(result_or.is_ok(), True)
+
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+
+    # Create client with registered methods
+    client2_with_listed_methods = rpc_ops.GrpcClient(
+        address=address, name="client2", list_registered_methods=True)
+
+    result_or = client2_with_listed_methods.assign_add(
+        variables.Variable(2, dtype=dtypes.int64))
+    self.assertAllEqual(result_or.is_ok(), True)
+
+    self.assertAllEqual(v, 6)
+
+    # Register new method after server started.
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError,
+        "All methods must be registered before starting the server"):
+      server.register("read_var", read_var)
+
+  def test_client_timeout(self):
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def add(a, b):
+      return math_ops.add(a, b)
+
+    server = rpc_ops.GrpcServer(address)
+
+    def start_server():
+      # Delay server start to simulate deadline exceeded for 1st RPC call
+      # response. Client waits till server is started, thus it can trigger
+      # deadline exceeded.
+      time.sleep(1)
+      server.register("add", add)
+      server.start()
+
+    t = threading.Thread(target=start_server)
+    t.start()
+
+    def ensure_server_is_ready(client):
+      server_ready = False
+      while not server_ready:
+        result_or = client.call(
+            "add", [constant_op.constant(20),
+                    constant_op.constant(30)])
+        if result_or.is_ok():
+          server_ready = True
+        else:
+          error_code, _ = result_or.get_error()
+          if error_code == errors.UNAVAILABLE:
+            server_ready = False
+          else:
+            server_ready = True
+      return
+
+    # Create client with list_registered_methods fails before server is started.
+    with self.assertRaises(errors.DeadlineExceededError):
+      rpc_ops.GrpcClient(
+          address,
+          name="client1",
+          list_registered_methods=True,
+          timeout_in_ms=1)
+
+    # Create same client again should succeed with
+    # list_registered_methods=False. Default timeout for client is 1 ms.
+    client = rpc_ops.GrpcClient(
+        address, name="client1", list_registered_methods=False, timeout_in_ms=1)
+
+    ensure_server_is_ready(client)
+    # Make explicit RPC call, the timeout of 1 ms should lead to
+    # deadline exceeded error.
+
+    result_or = client.call(
+        "add", [constant_op.constant(20),
+                constant_op.constant(30)],
+        timeout_in_ms=1)
+    self.assertAllEqual(result_or.is_ok(), False)
+    error_code, error_message = result_or.get_error()
+    self.assertAllEqual(error_code, errors.DEADLINE_EXCEEDED, error_message)
+
+    # Specifying reasonable timeout for call should succeed.
+    result_or = client.call(
+        "add", [constant_op.constant(20),
+                constant_op.constant(30)],
+        timeout_in_ms=5000)
+    self.assertAllEqual(result_or.is_ok(), True)
+    error_code, _ = result_or.get_error()
+
+    # Test timeouts for convenience methods
+
+    # Restart server again with delay to simulate deadline exceeded.
+    del server
+    server = rpc_ops.GrpcServer(address)
+    t = threading.Thread(target=start_server)
+    t.start()
+
+    # Client with no default timeout.
+    client = rpc_ops.GrpcClient(
+        address, name="client2", list_registered_methods=True)
+
+    # Succeeds with reasonable timeout.
+    result_or = client.add(
+        constant_op.constant(20), constant_op.constant(30), timeout_in_ms=5000)
+    self.assertAllEqual(result_or.is_ok(), True)
+
+  def test_async_call_op_wrapper(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign_add", assign_add)
+    server.register("read_var", read_var)
+    server.start()
+
+    client = rpc_ops.GrpcClient(address)
+
+    futures = []
+    for _ in range(10):
+      futures.append(
+          client.call("assign_add",
+                      [variables.Variable(2, dtype=dtypes.int64)]))
+
+    for f in futures:
+      f.is_ok()
+
+    result_or = client.call(
+        "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [20])
+
+  def test_rpc_call_op_in_tf_function(self):
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    server_resource.register("remote_fn", _remote_fn)
+
+    server_resource.start()
+    client = rpc_ops.GrpcClient(address=address, name="test_client")
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    @eager_def_function.function
+    def call_fn():
+      result_or = client.call(
+          args=[a, b],
+          method_name="remote_fn",
+          output_specs=[tensor_spec.TensorSpec([], dtypes.int32)])
+
+      self.assertAllEqual(True, result_or.is_ok())
+      result = result_or.get_value()
+      self.assertEqual(len(result), 1)  # Call returns a list(tensors)
+      # TODO(ishark): Shape for output tensor is unknown currently.
+      # Add attribute for capturing TensorSpec for output and enable
+      # check below:
+      # self.assertIsNotNone(result[0].shape.rank)
+      return result
+
+    self.assertAllEqual(call_fn(), [6])
+
+  def test_resource_deletion(self):
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server_handle = server._server_handle
+
+    # Test Future resource deletion
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    server.register("read_var", read_var)
+
+    server.start()
+    client = rpc_ops.GrpcClient(address)
+
+    client_handle = client._client_handle
+
+    # Check future resource deletion without calling get_value.
+    def _create_and_delete_rpc_future():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      return handle._status_or
+
+    @eager_def_function.function
+    def _create_and_delete_rpc_future_fn():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      return handle._status_or
+
+    for _ in range(2):
+      handle = _create_and_delete_rpc_future()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            handle, ignore_lookup_error=False)
+
+    for _ in range(2):
+      handle = _create_and_delete_rpc_future_fn()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            handle, ignore_lookup_error=False)
+
+    # Check future resource deletion with calling get_value.
+    def _create_and_delete_with_future():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      status_or_handle = handle._status_or
+      handle.get_value()
+      return status_or_handle
+
+    # Check future resource deletion with calling get_value with tf.function.
+    @eager_def_function.function
+    def _create_and_delete_with_future_fn():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      status_or_handle = handle._status_or
+      handle.get_value()
+      return status_or_handle
+
+    for _ in range(2):
+      resource_handle = _create_and_delete_with_future()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            resource_handle, ignore_lookup_error=False)
+
+    for _ in range(2):
+      resource_handle = _create_and_delete_with_future_fn()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            resource_handle, ignore_lookup_error=False)
+
+    # Test server client resource gets deleted.
+    del client
+    with self.assertRaises(errors.NotFoundError):
+      resource_variable_ops.destroy_resource_op(
+          client_handle, ignore_lookup_error=False)
+
+    # Test server server resource gets deleted.
+    del server
+    with self.assertRaises(errors.NotFoundError):
+      resource_variable_ops.destroy_resource_op(
+          server_handle, ignore_lookup_error=False)
+
+  def test_rpc_error(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign_add", assign_add)
+    server.register("read_var", read_var)
+    server.start()
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+
+    # confirm it works as expected when arguments are passed.
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    result_or = client.call(
+        "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [2])
+    result_or = client.assign_add(variables.Variable(2, dtype=dtypes.int64))
+    self.assertAllEqual(True, result_or.is_ok())
+
+    result_or = client.read_var()
+    self.assertAllEqual(True, result_or.is_ok())
+    self.assertAllEqual(result_or.get_value(), 4)
+
+    # Fails with invalid argument error when no arguments are passed.
+    result_or = client.call("assign_add")
+    self.assertAllEqual(result_or.is_ok(), False)
+    error_code, _ = result_or.get_error()
+    self.assertAllEqual(error_code, errors.INVALID_ARGUMENT)
+
+    del server
+    with self.assertRaises(errors.DeadlineExceededError):
+      _ = client.assign_add_blocking(
+          variables.Variable(2, dtype=dtypes.int64), timeout_in_ms=1)
+
+  def test_captured_inputs(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign_add", assign_add)
+    server.register("read_var", read_var)
+
+    server.start()
+
+    client = rpc_ops.GrpcClient(address)
+
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    result_or = client.call(
+        "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [4])
+
+  def test_register_method_twice(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign(a):
+      v.assign(a)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign", assign_add)
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "assign is already registered."):
+      # Reusing the same error name.
+      server.register("assign", assign)
+
+  def test_tf_function_register_without_input_signature(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function
+    def assign(a):
+      v.assign(a)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    with self.assertRaisesRegex(
+        ValueError, "Input signature not specified for the function."):
+      server.register("assign", assign)
+
+    # Register without input signature should work for functions without input
+    # args.
+    @eager_def_function.function
+    def read_var():
+      return v.value()
+
+    server.register("read_var", read_var)
+
+  def test_multi_device_resource(self):
+    elements = np.random.randint(100, size=[200])
+
+    with ops.device("/device:CPU:1"):
+      queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
+
+    @eager_def_function.function()
+    def populate_queue():
+      queue.enqueue_many(elements)
+      queue.close()
+
+    with ops.device("/device:CPU:0"):
+      port = portpicker.pick_unused_port()
+      address = "localhost:{}".format(port)
+      server = rpc_ops.GrpcServer(address)
+      server.register("populate_queue", populate_queue)
+      server.start()
+
+      client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+      client.populate_queue()
+
+    for e in elements:
+      self.assertAllEqual(e, queue.dequeue())
+
+  def test_queue_resource(self):
+    elements = np.random.randint(100, size=[200])
+    queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
+
+    @eager_def_function.function()
+    def populate_queue():
+      queue.enqueue_many(elements)
+      queue.close()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("populate_queue", populate_queue)
+    server.start()
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+    client.populate_queue()
+
+    for e in elements:
+      self.assertAllEqual(e, queue.dequeue())
+
+  def test_multi_device_resource_cpu(self):
+    with ops.device("/device:cpu:1"):
+      v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    with ops.device("/device:CPU:0"):
+      port = portpicker.pick_unused_port()
+      address = "localhost:{}".format(port)
+      server = rpc_ops.GrpcServer(address)
+      server.register("assign_add", assign_add)
+      server.start()
+
+      client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+      result_or = client.assign_add(variables.Variable(2, dtype=dtypes.int64))
+      self.assertAllEqual(result_or.is_ok(), True)
+
+    self.assertAllEqual(v, 2)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()