chore: import upstream snapshot with attribution

tensorflow/python/distribute/parallel_device/BUILD

@@ -0,0 +1,76 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow/core/platform:distribute.bzl", "distribute_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+# Pybind rules must live in tensorflow/python due to header rule visibility.
+exports_files(
+    ["pywrap_parallel_device.cc"],
+    visibility = ["//tensorflow/python:__pkg__"],
+)
+
+py_library(
+    name = "parallel_device",
+    srcs = ["parallel_device.py"],
+    strict_deps = True,
+    visibility = [
+        "//tensorflow:internal",
+        "//third_party/py/tensorflow_numerics/extensions:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/python:_pywrap_parallel_device",
+        "//tensorflow/python/distribute:device_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/tpu/ops",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:variable_utils",
+    ],
+)
+
+distribute_py_strict_test(
+    name = "parallel_device_test",
+    srcs = ["parallel_device_test.py"],
+    # We don't only use TPU Runtime on parallel device, we use c_api_tfrt instead.
+    disable_tpu_use_tfrt = True,
+    tags = [
+        # Dependencies aren't otherwise included in the pip package yet.
+        "no_pip",
+        "notpu",  # TODO(b/210159281)
+    ],
+    deps = [
+        ":parallel_device",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_management",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:collective_ops",
+        "//tensorflow/python/ops:control_flow_switch_case",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:stateful_random_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/util:nest",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/distribute/parallel_device/parallel_device.py

@@ -0,0 +1,232 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility for eagerly executing operations in parallel on multiple devices."""
+
+import threading
+import weakref
+
+from tensorflow.python import _pywrap_parallel_device
+from tensorflow.python.distribute import device_util
+from tensorflow.python.eager import context
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.tpu.ops import tpu_ops
+from tensorflow.python.util import nest
+from tensorflow.python.util import variable_utils
+
+_next_device_number = 0
+_next_device_number_lock = threading.Lock()
+
+_all_parallel_devices = weakref.WeakValueDictionary()
+
+
+def unpack(tensor):
+  """Finds `tensor`'s parallel device and unpacks its components."""
+  parallel_device = _all_parallel_devices.get(tensor.device, None)
+  if parallel_device is None:
+    raise ValueError("{} is not a parallel device".format(tensor.device))
+  return parallel_device.unpack(tensor)
+
+
+# TODO(allenl): Expand this docstring once things like getting components on and
+# off the device are stable.
+#
+# TODO(allenl): Make multi-client work; we need an offset for device IDs, and an
+# indication of how many other devices there are total for collectives which
+# don't have a number of participants hard-coded in their attributes.
+class ParallelDevice(object):
+  """A device which executes operations in parallel."""
+
+  def __init__(self, components):
+    """Creates a device which executes operations in parallel on `components`.
+
+    Args:
+      components: A list of device names. Each operation executed on the
+        returned device executes on these component devices.
+
+    Returns:
+      A string with the name of the newly created device.
+    """
+    global _next_device_number, _next_device_number_lock
+    self.components = tuple(device_util.canonicalize(d) for d in components)
+    if not self.components:
+      raise ValueError("ParallelDevice requires at least one component.")
+    ctx = context.context()
+    with _next_device_number_lock:
+      # TODO(allenl): Better names for parallel devices (right now "CUSTOM" is
+      # special-cased).
+      self._name = "{}/device:CUSTOM:{}".format(ctx.host_address_space(),
+                                                _next_device_number)
+      _next_device_number += 1
+    device, device_info = _pywrap_parallel_device.GetParallelDeviceCapsules(
+        self._name, self.components)
+    context.register_custom_device(device, self._name, device_info)
+    self._device_ids = None
+    self._device_scope = None
+    _all_parallel_devices[self._name] = self
+
+  def _pack_tensor(self, *tensors):
+    """Helper to pack plain-old-tensors, not structures or composites."""
+    for tensor in tensors:
+      if not isinstance(
+          tensor,
+          (
+              tensor_lib.Tensor,
+              composite_tensor.CompositeTensor,
+              variables.Variable,
+          ),
+      ):
+        raise ValueError(
+            ("Every component to pack onto the ParallelDevice must already be "
+             "a tensor, got {}. Consider running `tf.constant` or "
+             "`tf.convert_to_tensor` first on literal values.")
+            .format(tensors))
+    with ops.device(self._name):
+      return tpu_ops.tpu_replicated_input(inputs=tensors)
+
+  def pack(self, tensors):
+    """Create a tensor on the parallel device from a sequence of tensors.
+
+    Args:
+      tensors: A list of tensors, one per device in `self.components`. The list
+        can contain composite tensors and nests (lists, dicts, etc. supported by
+        `tf.nest`) with the same structure for each device, but every component
+        of nests must already be a `tf.Tensor` or composite. Passing
+        `tf.Variable` objects reads their value, it does not share a mutable
+        reference between the packed and unpacked forms.
+
+    Returns:
+      A tensor placed on the ParallelDevice. For nested structures, returns a
+      single structure containing tensors placed on the ParallelDevice (same
+      structure as each component of `tensors`).
+
+    Raises:
+      ValueError: If the length of `tensors` does not match the number of
+        component devices, or if there are non-tensor inputs.
+
+    """
+    self._assert_eager()
+    if len(tensors) != len(self.components):
+      raise ValueError(
+          ("Creating a parallel tensor requires one tensor per component. "
+           "Got {} but was expecting {}.")
+          .format(len(tensors), len(self.components)))
+    with ops.device(None):
+      # Explicitly read variable values. This can not be done on the parallel
+      # device since the tensors are to be packed.
+      tensors = variable_utils.convert_variables_to_tensors(tensors)
+    return nest.map_structure(self._pack_tensor, *tensors,
+                              expand_composites=True)
+
+  def _unpack_tensor(self, parallel_tensor):
+    """Helper to unpack a single tensor."""
+    if not isinstance(
+        parallel_tensor,
+        (
+            tensor_lib.Tensor,
+            composite_tensor.CompositeTensor,
+            variables.Variable,
+        ),
+    ):
+      raise ValueError("Expected a tensor, got {}.".format(parallel_tensor))
+    with ops.device(self._name):
+      return tpu_ops.tpu_replicated_output(
+          parallel_tensor, num_replicas=len(self.components))
+
+  def unpack(self, parallel_tensor):
+    """Unpack a parallel tensor into its components.
+
+    Args:
+      parallel_tensor: A tensor, composite tensor, or `tf.nest` of such placed
+        on the ParallelDevice. Passing `tf.Variable` objects reads their value,
+        it does not share a mutable reference between the packed and unpacked
+        forms.
+
+    Returns:
+      A list with the same length as `self.components` each with the same
+      structure as `parallel_tensor`, containing component tensors.
+
+    """
+    self._assert_eager()
+    unpacked_components = [[] for _ in range(len(self.components))]
+    with ops.device(self._name):
+      parallel_tensor = variable_utils.convert_variables_to_tensors(
+          parallel_tensor)
+    for tensor in nest.flatten(parallel_tensor, expand_composites=True):
+      for accumulator, unpacked_tensor in zip(
+          unpacked_components, self._unpack_tensor(tensor)):
+        accumulator.append(unpacked_tensor)
+    return [nest.pack_sequence_as(parallel_tensor, unpacked,
+                                  expand_composites=True)
+            for unpacked in unpacked_components]
+
+  @property
+  def device_ids(self):
+    """A parallel tensor with scalar integers numbering component devices.
+
+    Each device ID is placed on its corresponding device, in the same order as
+    the `components` constructor argument.
+
+    Returns:
+      A parallel tensor containing 0 on the first device, 1 on the second, etc.
+    """
+    if self._device_ids is None:
+      # device_ids may be called from inside a tf.function, in which case the
+      # function captures the eager tensor. We can't pack tensors in a function
+      # at the moment, and even if we could we don't want to hold on to a
+      # symbolic tensor, so we need to init_scope out of the function
+      # temporarily.
+      with ops.init_scope():
+        # TODO(allenl): Functions which capture eager device ID tensors won't be
+        # saveable in SavedModels. Ideally we'd run a DeviceID op every time
+        # device IDs are required, with functions using the op in their bodies
+        # but not hard-coding a fixed number of devices (so they can be re-used
+        # with a different replica count).
+        device_ids_list = []
+        for index, device in enumerate(self.components):
+          with ops.device(device):
+            # The identity op ensures each device ID tensor is placed on its
+            # device.
+            device_ids_list.append(
+                array_ops.identity(constant_op.constant(index)))
+        self._device_ids = self.pack(device_ids_list)
+
+    return self._device_ids
+
+  def _assert_eager(self):
+    """Verifies that tracing is not active."""
+    if not context.executing_eagerly():
+      raise NotImplementedError(
+          "ParallelDevice is currently not supported inside `tf.function`. It "
+          "can however run calls to a `tf.function` in parallel:\n\n"
+          "with ParallelDevice() as p:\n  f()")
+
+  def __enter__(self):
+    """Runs ops in parallel, makes variables which save independent buffers."""
+    if self._device_scope is not None:
+      raise AssertionError(
+          "Re-entered a ParallelDevice scope without first exiting it.")
+    self._assert_eager()
+    self._device_scope = ops.device(self._name)
+    self._device_scope.__enter__()
+    return self
+
+  def __exit__(self, typ, exc, tb):
+    self._device_scope.__exit__(typ, exc, tb)
+    self._device_scope = None

tensorflow/python/distribute/parallel_device/parallel_device_test.py

@@ -0,0 +1,622 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import os
+import threading
+
+from absl.testing import parameterized
+from tensorflow.python.checkpoint import checkpoint as tracking
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.distribute.parallel_device import parallel_device
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+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 errors
+from tensorflow.python.framework import ops
+from tensorflow.python.module import module
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.ops import control_flow_switch_case
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import stateful_random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import save
+from tensorflow.python.util import nest
+
+# When running collectives asynchronously, we need to give each parallel device
+# execution a unique ID so the collectives don't interfere. Since the op is
+# replicated with group/instance key intact, the replicated nodes will
+# communicate.
+# TODO(allenl): Switch to using a collective manager.
+_COUNTER_LOCK = threading.Lock()
+_COUNTER = 100
+
+
+def _collective_reduce(inputs, operation, num_replicas):
+
+  def _reduce_tensor(tensor):
+    with _COUNTER_LOCK:
+      global _COUNTER
+      keys = _COUNTER
+      _COUNTER += 1
+    return collective_ops.all_reduce_v2(
+        t=tensor,
+        group_size=num_replicas,
+        merge_op=operation,
+        group_key=keys,
+        instance_key=keys)
+
+  return nest.map_structure(_reduce_tensor, inputs)
+
+
+def _collective_sum(inputs, num_replicas):
+  return _collective_reduce(
+      inputs=inputs, operation="Add", num_replicas=num_replicas)
+
+
+class _Dense(module.Module):
+
+  def __init__(self, output_size):
+    self.output_size = output_size
+    self.kernel = None
+    self.bias = None
+
+  def __call__(self, x):
+    if self.kernel is None:
+      self.kernel = variables.Variable(
+          array_ops.ones(
+              array_ops_stack.stack([self.output_size,
+                                     array_ops.shape(x)[-1]])))
+      self.bias = variables.Variable(array_ops.ones([self.output_size]))
+    return math_ops.matmul(x, self.kernel, transpose_b=True) + self.bias
+
+
+class _VirtualDeviceTestCase(test.TestCase):
+
+  def setUp(self):
+    super(_VirtualDeviceTestCase, self).setUp()
+    ctx = context.context()
+    if ctx.list_physical_devices("TPU"):
+      self.device_type = "TPU"
+      tpu_cluster_resolver.initialize_tpu_system()
+    elif ctx.list_physical_devices("GPU"):
+      self.device_type = "GPU"
+      gpus = ctx.list_physical_devices(self.device_type)
+      ctx.set_logical_device_configuration(gpus[0], [
+          context.LogicalDeviceConfiguration(memory_limit=100),
+          context.LogicalDeviceConfiguration(memory_limit=100),
+      ])
+    else:
+      self.device_type = "CPU"
+      cpus = ctx.list_physical_devices("CPU")
+      ctx.set_logical_device_configuration(cpus[0], [
+          context.LogicalDeviceConfiguration(),
+          context.LogicalDeviceConfiguration(),
+      ])
+
+    self.device = parallel_device.ParallelDevice(components=[
+        "/job:localhost/device:{}:0".format(self.device_type),
+        self.device_type + ":1"
+    ])
+    self.assertIn(self.device_type + ":0", self.device.components[0])
+    self.assertIn(self.device_type + ":1", self.device.components[1])
+
+
+class ParallelDeviceTests(_VirtualDeviceTestCase, parameterized.TestCase):
+
+  def test_register_parallel_device(self):
+    with self.device:
+      c = constant_op.constant(1.)
+      d = constant_op.constant(2.)
+      e = c + d
+      outputs = self.device.unpack(e)
+    self.assertAllClose([3., 3.], outputs)
+
+    self.assertIn(self.device.components[0], outputs[0].backing_device)
+    self.assertIn(self.device.components[1], outputs[1].backing_device)
+
+  def test_no_implicit_copyon(self):
+    a1 = constant_op.constant(1.)
+    a2 = constant_op.constant(2.)
+
+    with self.device:
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          "First pack non-parallel tensors for each device"):
+        a1 + a2  # pylint:disable=pointless-statement
+
+  def test_error_message_length(self):
+    x = array_ops.ones([3, 3, 3, 3, 3, 3])
+
+    with self.device:
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          r"TensorHandle\((.|\n){1,150}\[...\], shape="):
+        array_ops.identity(x)
+
+  def test_one_replica_eager_control_flow(self):
+    device = parallel_device.ParallelDevice(components=[
+        "/job:localhost/device:{}:0".format(self.device_type),
+    ])
+    x = constant_op.constant([2, 3, 4])
+    with device:
+      x = device.pack([x])
+      if math_ops.reduce_any(math_ops.equal(x, constant_op.constant(4))):
+        y = constant_op.constant(1)
+      else:
+        y = constant_op.constant(2)
+    self.assertAllEqual([1], device.unpack(y))
+
+  @parameterized.named_parameters(
+      ("variable", variables.Variable),
+      ("tensor", lambda x: x))
+  def test_string_representation(self, transform):
+    x = self.device.pack(
+        [constant_op.constant([5., 6.]),
+         constant_op.constant([6., 7.])])
+    with self.device:
+      x = transform(x)
+    parallel_str = str(x)
+    self.assertIn("5", parallel_str)
+    self.assertIn("7", parallel_str)
+    self.assertIn(self.device_type + ":0", parallel_str)
+    self.assertIn(self.device_type + ":1", parallel_str)
+    parallel_repr = repr(x)
+    self.assertIn("5", parallel_repr)
+    self.assertIn("7", parallel_repr)
+    self.assertIn(self.device_type + ":0", parallel_repr)
+    self.assertIn(self.device_type + ":1", parallel_repr)
+
+  def test_device_id(self):
+    device_ids = self.device.unpack(self.device.device_ids)
+    self.assertAllClose([0, 1], device_ids)
+    # TODO(allenl): Should device IDs be int64 so they can be placed on GPUs?
+    # Currently backing_device is CPU.
+    self.assertIn(self.device.components[0], device_ids[0].device)
+    self.assertIn(self.device.components[1], device_ids[1].device)
+
+  def test_zeros(self):
+    with self.device:
+      x = array_ops.zeros([array_ops.identity(constant_op.constant(10))])
+    for component in self.device.unpack(x):
+      self.assertAllClose([0.] * 10, component)
+
+  def test_generator(self):
+    with self.device:
+      g_same = stateful_random_ops.Generator.from_seed(0)
+      g_different = stateful_random_ops.Generator.from_seed(
+          self.device.device_ids)
+      same = g_same.normal([10])
+      different = g_different.normal([10])
+    same_unpacked = self.device.unpack(same)
+    different_unpacked = self.device.unpack(different)
+    for same_component, different_component in zip(same_unpacked[1:],
+                                                   different_unpacked[1:]):
+      self.assertAllClose(same_component, same_unpacked[0])
+      self.assertNotAllClose(different_component, different_unpacked[0])
+
+  def test_collective_reduce(self):
+    x = self.device.pack(
+        [constant_op.constant(-1.5),
+         constant_op.constant(3.5)])
+    with self.device:
+      reduced = _collective_sum(x, num_replicas=2)
+      outputs = self.device.unpack(reduced)
+    self.assertAllClose([2., 2.], outputs)
+    self.assertIn(self.device.components[0], outputs[0].backing_device)
+    self.assertIn(self.device.components[1], outputs[1].backing_device)
+
+  def test_collective_reduce_in_function(self):
+    x = self.device.pack(
+        [constant_op.constant(-1.5),
+         constant_op.constant(3.5)])
+    with self.device:
+
+      @def_function.function
+      def reduce(t):
+        return _collective_sum(t, num_replicas=2)
+
+      reduced = reduce(x)
+      outputs = self.device.unpack(reduced)
+    self.assertAllClose([2., 2.], outputs)
+    self.assertIn(self.device.components[0], outputs[0].backing_device)
+    self.assertIn(self.device.components[1], outputs[1].backing_device)
+
+  def test_collective_reduce_async_scope(self):
+    # Note that ops on the parallel device currently don't execute
+    # asynchronously. The test is just that we don't get deadlocks.
+    x = self.device.pack(
+        [constant_op.constant(-1.5),
+         constant_op.constant(3.5)])
+    with context.async_scope(), self.device:
+      reduced = _collective_sum(x, num_replicas=2)
+      outputs = self.device.unpack(reduced)
+    self.assertAllClose([2., 2.], outputs)
+    self.assertIn(self.device.components[0], outputs[0].backing_device)
+    self.assertIn(self.device.components[1], outputs[1].backing_device)
+
+  def test_collective_reduce_async_context(self):
+    previous = config.get_synchronous_execution()
+    try:
+      context._reset_context()
+      config.set_synchronous_execution(False)
+      self.setUp()
+      # Note that ops on the parallel device currently don't execute
+      # asynchronously. The test is just that we don't get deadlocks.
+      x = self.device.pack(
+          [constant_op.constant(-1.5),
+           constant_op.constant(3.5)])
+      with self.device:
+        reduced = _collective_sum(x, num_replicas=2)
+        outputs = self.device.unpack(reduced)
+      self.assertAllClose([2., 2.], outputs)
+      self.assertIn(self.device.components[0], outputs[0].backing_device)
+      self.assertIn(self.device.components[1], outputs[1].backing_device)
+    finally:
+      context._reset_context()
+      config.set_synchronous_execution(previous)
+
+  def test_collective_broadcast_in_function(self):
+    if self.device_type == "TPU":
+      self.skipTest("ParallelDevice broadcast collectives on TPUs need work")
+
+    @def_function.function
+    def broadcast_send_recv(device_id):
+      c = constant_op.constant([2])
+
+      @def_function.function
+      def send():
+        s0 = collective_ops.broadcast_send(
+            c * 3, c.shape, c.dtype, group_size=2, group_key=1, instance_key=1)
+        with ops.control_dependencies([s0.op]):
+          return array_ops.identity(c)
+
+      @def_function.function
+      def recv():
+        r0 = collective_ops.broadcast_recv(
+            c.shape, c.dtype, group_size=2, group_key=1, instance_key=1)
+        return r0
+
+      return control_flow_switch_case.switch_case(
+          device_id, branch_fns={
+              0: send,
+              1: recv
+          })
+
+    with self.device:
+      result = broadcast_send_recv(self.device.device_ids)
+    self.assertAllClose([[2], [6]], self.device.unpack(result))
+
+  def test_use_in_graph_error_is_informative(self):
+    @def_function.function
+    def uses_parallel():
+      with self.device:
+        return self.device.unpack(array_ops.ones([]))
+
+    with self.assertRaisesRegex(NotImplementedError, "inside `tf.function`"):
+      uses_parallel()
+
+  def test_checkpointing(self):
+    self.skipTest("b/216201668: revisit parallel device and checkpointing.")
+
+    prefix = os.path.join(self.get_temp_dir(), "ckpt")
+    different_values = self.device.pack(
+        [constant_op.constant(-1.),
+         constant_op.constant(3.)])
+    with self.device:
+      v = variables.Variable(different_values)
+      checkpoint = tracking.Checkpoint(v=v)
+    save_path = checkpoint.save(prefix)
+    with self.device:
+      v.assign(constant_op.constant(0.))
+    checkpoint.restore(save_path).assert_consumed()
+    with self.device:
+      outputs = self.device.unpack(v)
+    self.assertAllClose([-1., 3.], outputs)
+
+    with self.device:
+      restore_on_create = tracking.Checkpoint()
+      restore_on_create.restore(save_path)
+      restore_on_create.v = variables.Variable(0.)
+      outputs = self.device.unpack(restore_on_create.v)
+    self.assertAllClose([-1., 3.], outputs)
+
+    # Changing the number of devices / restoring into a single-device copy is OK
+    single_device = tracking.Checkpoint(v=variables.Variable(0.))
+    status = single_device.restore(save_path)
+    status.assert_existing_objects_matched()
+    self.assertAllClose(-1., single_device.v)
+    with self.assertRaisesRegex(AssertionError, "parallel_component_1"):
+      # There are parts of the variable that aren't restored into a
+      # single-device copy.
+      status.assert_consumed()
+
+  def test_pack_composite(self):
+    if self.device_type != "CPU":
+      self.skipTest("Iterator GetNext doesn't work on accelerators.")
+    datasets = [
+        dataset_ops.Dataset.from_tensor_slices(
+            [i + 1, (i + 1) * 2, (i + 1) * 3])
+        for i in range(len(self.device.components))]
+    parallel_dataset = self.device.pack(datasets)
+    with self.device:
+      iterator = iter(parallel_dataset)
+      parallel_sample = next(iterator)
+    component_iterators = self.device.unpack(iterator)
+    self.assertEqual(2, next(component_iterators[0]).numpy())
+    self.assertEqual(1, self.device.unpack(parallel_sample)[0].numpy())
+    self.assertEqual(4, next(component_iterators[1]).numpy())
+    self.assertEqual(2, self.device.unpack(parallel_sample)[1].numpy())
+
+  def test_pack_structure(self):
+    x_parts = [{"a": constant_op.constant(float(i))}
+               for i in range(len(self.device.components))]
+    x = self.device.pack(x_parts)
+    self.assertAllClose([{"a": 0.}, {"a": 1.}], self.device.unpack(x))
+
+  def test_pack_variable_value(self):
+    x_parts = [variables.Variable(i)
+               for i in range(len(self.device.components))]
+    x = self.device.pack(x_parts)
+    with self.device:
+      x1 = self.device.pack(x_parts)
+    for v in x_parts:
+      v.assign(-10)  # Mutating the variable does not affect previous reads.
+    self.assertAllClose([0, 1], self.device.unpack(x))
+    self.assertAllClose([0, 1], self.device.unpack(x1))
+
+  def test_unpack_variable_value(self):
+    x_parts = [constant_op.constant(i)
+               for i in range(len(self.device.components))]
+    x = self.device.pack(x_parts)
+    with self.device:
+      v = variables.Variable(x)
+      v_unpacked = self.device.unpack(v)
+      v.assign(-10)  # Mutating the variable does not affect previous reads.
+    self.assertAllClose([0, 1], v_unpacked)
+
+  def test_saved_model(self):
+    self.skipTest("b/216201668: revisit parallel device and saved model")
+
+    different_values = self.device.pack(
+        [constant_op.constant(-1.),
+         constant_op.constant(3.)])
+    with self.device:
+      m = module.Module()
+      m.v = variables.Variable(different_values)
+      m.f = def_function.function(lambda: m.v * 2.)
+      self.assertAllClose([-2., 6.], self.device.unpack(m.f()))
+    saved_model_path = os.path.join(self.get_temp_dir(), "saved_model")
+    save.save(m, saved_model_path)
+
+    context._reset_context()
+    self.setUp()
+
+    single_device_loaded = load.load(saved_model_path)
+    self.assertAllClose(-2., single_device_loaded.f())
+    assign_value = self.device.pack(
+        [constant_op.constant(.1), constant_op.constant(.2)])
+    with self.device:
+      parallel_loaded = load.load(saved_model_path)
+      self.assertAllClose([-2., 6.], self.device.unpack(parallel_loaded.f()))
+      self.assertAllClose([-1., 3.], self.device.unpack(parallel_loaded.v))
+      parallel_loaded.v.assign(assign_value)
+      self.assertAllClose([.2, .4], self.device.unpack(parallel_loaded.f()))
+
+  def _assert_close_to_non_parallel(self, computation):
+    """Asserts that replication of `computation` works and is equivalent."""
+    with self.device:
+      parallel_result = computation()
+    non_parallel_result = computation()
+    # The computations should have the same number and structure of Tensor
+    # objects, even though the tensors themselves will be on different devices
+    # and represent different numbers of values.
+    nest.assert_same_structure(parallel_result, non_parallel_result)
+    non_parallel_flat = nest.flatten(non_parallel_result)
+    parallel_flat = nest.flatten(parallel_result)
+    self.assertGreater(len(parallel_flat), 0)
+    for non_parallel, parallel in zip(non_parallel_flat, parallel_flat):
+      self.assertEqual(self.device._name, parallel.device)
+      self.assertNotEqual(self.device._name, non_parallel.device)
+      for parallel_component in self.device.unpack(parallel):
+        self.assertAllClose(non_parallel, parallel_component)
+
+  def test_capturing(self):
+    with self.device:
+      x = constant_op.constant([1., 2.])
+      x = array_ops.identity(x)
+
+      @def_function.function
+      def f(y):
+        return x + y
+
+      y = array_ops.ones([2])
+      parallel_result = f(y)
+    self.assertAllClose([[2., 3.]] * 2, self.device.unpack(parallel_result))
+
+  def test_euclidean_norm(self):
+    def _test_fn():
+      with backprop.GradientTape() as tape:
+        x = array_ops.ones([5, 5])
+        tape.watch(x)
+        y = math_ops.reduce_euclidean_norm(x, axis=constant_op.constant(1))
+      return y, tape.gradient(y, x)
+    self._assert_close_to_non_parallel(_test_fn)
+
+  def test_reduce_sum(self):
+    def _test_fn():
+      with backprop.GradientTape() as tape:
+        x = array_ops.ones([5, 5])
+        tape.watch(x)
+        y = math_ops.reduce_sum(x, axis=constant_op.constant(1))
+      return y, tape.gradient(y, x)
+    self._assert_close_to_non_parallel(_test_fn)
+
+  def test_variable_created_in_function(self):
+    captured_value = constant_op.constant(2.)
+
+    class M(module.Module):
+
+      def __init__(self):
+        self.v = None
+        self.w = None
+        self.x = None
+        self.z = None
+
+      @def_function.function(autograph=False)
+      def __call__(self, x):
+        if self.v is None:
+          with ops.init_scope():
+            initial_value = constant_op.constant(2.)
+            self.z = variables.Variable(initial_value)
+          self.x = variables.Variable(captured_value)
+          self.w = variables.Variable(lambda: constant_op.constant(2.))
+          self.v = variables.Variable(constant_op.constant(2.))
+        return x * self.v * self.w * self.x * self.z
+
+    with self.device:
+      m = M()
+      packed_outputs = m(array_ops.ones([]))
+      outputs = self.device.unpack(packed_outputs)
+    self.assertAllClose([16., 16.], outputs)
+
+  def test_different_shapes(self):
+    x = self.device.pack(
+        [constant_op.constant([1., 2.]),
+         constant_op.constant([5.])])
+    with self.device:
+      y = x * 2.
+    self.assertEqual([None], y.shape.as_list())
+    self.assertAllClose([[2., 4.], [10.]], self.device.unpack(y))
+
+    different_axes = self.device.pack(
+        [constant_op.constant([1., 2.]),
+         constant_op.constant([[5.]])])
+    with self.assertRaisesRegex(Exception,
+                                "components do not all have the same rank"):
+      different_axes.shape  # pylint: disable=pointless-statement
+
+
+class LayerTests(_VirtualDeviceTestCase):
+
+  def test_layer_forward(self):
+    with self.device:
+      layer = _Dense(5)
+      x = constant_op.constant([[2.]])
+      y = layer(x)
+      outputs = self.device.unpack(y)
+    self.assertAllClose([[3.] * 5], outputs[0])
+    self.assertAllClose([[3.] * 5], outputs[1])
+    self.assertIn(self.device.components[0], outputs[0].backing_device)
+    self.assertIn(self.device.components[1], outputs[1].backing_device)
+
+    # With different Layer inputs we get different outputs
+    x = self.device.pack(
+        [constant_op.constant([[-0.5]]),
+         constant_op.constant([[0.5]])])
+    with self.device:
+      y = layer(x)
+      outputs = self.device.unpack(y)
+    self.assertGreater(
+        math_ops.reduce_max(math_ops.abs(outputs[0] - outputs[1])), 1e-5)
+    self.assertIn(self.device.components[0], outputs[0].backing_device)
+    self.assertIn(self.device.components[1], outputs[1].backing_device)
+
+  def test_layer_sync_training(self):
+    x = self.device.pack(
+        [constant_op.constant([[-0.5]]),
+         constant_op.constant([[0.5]])])
+    with self.device:
+      layer = _Dense(5)
+
+      with backprop.GradientTape() as tape:
+        y = layer(x)
+        loss = (y - math_ops.range(5.))**2.
+      parameters = layer.trainable_variables
+      unreduced_gradients = tape.gradient(loss, parameters)
+      reduced_gradients = _collective_sum(unreduced_gradients, num_replicas=2)
+      for grad, param in zip(reduced_gradients, parameters):
+        param.assign_sub(0.01 * grad)
+    final_kernels = self.device.unpack(layer.kernel)
+    self.assertAllClose(final_kernels[0], final_kernels[1])
+    final_bias = self.device.unpack(layer.bias)
+    expected_bias = (1. - 0.01 * 2. * (1. + .5 - math_ops.range(5.)) -
+                     0.01 * 2. * (1. - .5 - math_ops.range(5.)))
+    self.assertAllClose(expected_bias, final_bias[0], rtol=1e-4, atol=1e-4)
+    self.assertAllClose(expected_bias, final_bias[1], rtol=1e-4, atol=1e-4)
+    self.assertIn(self.device.components[0], final_kernels[0].backing_device)
+    self.assertIn(self.device.components[1], final_kernels[1].backing_device)
+
+  def test_layer_divergent_buffer_training(self):
+    x = self.device.pack(
+        [constant_op.constant([[-0.5]]),
+         constant_op.constant([[0.5]])])
+    with self.device:
+      layer = _Dense(5)
+
+      with backprop.GradientTape() as tape:
+        y = layer(x)
+        loss = (y - math_ops.range(5.))**2.
+      parameters = layer.trainable_variables
+      unreduced_gradients = tape.gradient(loss, parameters)
+      for grad, param in zip(unreduced_gradients, parameters):
+        param.assign_sub(0.01 * grad)
+    final_kernels = self.device.unpack(layer.kernel)
+    self.assertNotAllClose(final_kernels[0], final_kernels[1])
+    final_bias = self.device.unpack(layer.bias)
+    self.assertAllClose(1. - 0.01 * 2. * (1. - .5 - math_ops.range(5.)),
+                        final_bias[0])
+    self.assertAllClose(1. - 0.01 * 2. * (1. + .5 - math_ops.range(5.)),
+                        final_bias[1])
+    self.assertIn(self.device.components[0], final_kernels[0].backing_device)
+    self.assertIn(self.device.components[1], final_kernels[1].backing_device)
+
+  def test_training_loop(self):
+    self.skipTest("b/216201668: revisit parallel device and checkpointing")
+    for _ in range(5):
+      layer = _Dense(5)
+      checkpoint = tracking.Checkpoint(layer=layer)
+      manager = checkpoint_management.CheckpointManager(
+          checkpoint, directory=self.get_temp_dir(), max_to_keep=5)
+      manager.restore_or_initialize()
+
+      for _ in range(10):
+        x = self.device.pack(
+            [constant_op.constant([[-0.5]]),
+             constant_op.constant([[0.5]])])
+        with self.device:
+          with backprop.GradientTape() as tape:
+            y = layer(x)
+            loss = (y - math_ops.range(5.))**2.
+          parameters = layer.trainable_variables
+          unreduced_gradients = tape.gradient(loss, parameters)
+          reduced_gradients = _collective_sum(
+              unreduced_gradients, num_replicas=len(self.device.components))
+          for grad, param in zip(reduced_gradients, parameters):
+            param.assign_sub(0.01 * grad)
+
+        manager.save()
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/distribute/parallel_device/pywrap_parallel_device.cc

@@ -0,0 +1,68 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+#include <vector>
+
+#include "Python.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/eager/parallel_device/parallel_device.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
+
+namespace py = pybind11;
+
+void CallDelete_Device(PyObject* capsule) {
+  delete reinterpret_cast<TFE_CustomDevice*>(
+      PyCapsule_GetPointer(capsule, "TFE_CustomDevice"));
+}
+
+void CallDelete_DeviceInfo(PyObject* capsule) {
+  void (*destructor)(void*) =
+      reinterpret_cast<void (*)(void*)>(PyCapsule_GetContext(capsule));
+  destructor(PyCapsule_GetPointer(capsule, "TFE_CustomDevice_DeviceInfo"));
+}
+
+PYBIND11_MODULE(_pywrap_parallel_device, m) {
+  m.def("GetParallelDeviceCapsules",
+        [](const char* name, std::vector<std::string> underlying_devices) {
+          std::vector<const char*> underlying_devices_c;
+          underlying_devices_c.reserve(underlying_devices.size());
+          for (const std::string& element : underlying_devices) {
+            underlying_devices_c.push_back(element.c_str());
+          }
+          // `device` is owned by `device_capsule`.
+          TFE_CustomDevice* device = new TFE_CustomDevice;
+          tensorflow::Safe_PyObjectPtr device_capsule(
+              PyCapsule_New(device, "TFE_CustomDevice", &CallDelete_Device));
+          void* device_info;
+          tensorflow::parallel_device::AllocateParallelDevice(
+              name, underlying_devices_c.data(), underlying_devices_c.size(),
+              device, &device_info);
+          if (PyErr_Occurred()) throw py::error_already_set();
+          tensorflow::Safe_PyObjectPtr device_info_capsule(
+              PyCapsule_New(device_info, "TFE_CustomDevice_DeviceInfo",
+                            &CallDelete_DeviceInfo));
+          if (PyErr_Occurred()) throw py::error_already_set();
+          // The PyCapsule destructor needs a pointer to the destructor for
+          // DeviceInfo.
+          PyCapsule_SetContext(device_info_capsule.get(),
+                               reinterpret_cast<void*>(device->delete_device));
+          return tensorflow::PyoOrThrow(
+              PyTuple_Pack(2, device_capsule.get(), device_info_capsule.get()));
+        });
+}