cross_device_ops_test.py

# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for CrossDeviceOps."""

import collections
import os
import threading
import time

from absl.testing import parameterized

from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import tensorflow_server_pb2
from tensorflow.python.distribute import cluster_resolver as cluster_resolver_lib
from tensorflow.python.distribute import collective_util
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
from tensorflow.python.distribute import cross_device_utils
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import multi_process_runner
from tensorflow.python.distribute import multi_worker_test_base
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import test_util
from tensorflow.python.distribute import values as value_lib
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import indexed_slices
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import collective_ops
from tensorflow.python.ops import cond
from tensorflow.python.ops import math_ops
from tensorflow.python.util import nest


try:
  import dill  # pylint:disable=g-import-not-at-top

  _REGISTER_DECORATOR = dill.register
except ImportError:
  _REGISTER_DECORATOR = lambda fn, *_: fn

CollectiveReplicaLauncher = cross_device_utils.CollectiveReplicaLauncher
CommunicationImplementation = collective_util.CommunicationImplementation
ReduceOp = reduce_util.ReduceOp
IndexedSlicesValue = indexed_slices.IndexedSlicesValue
IndexedSlices = indexed_slices.IndexedSlices


def make_per_replica_value(value, devices):
  """Creates a `PerReplica` object whose values reside in `devices`.

  Args:
    value: a tensor-convertible value or a `IndexedSlicesValue`, or a callable
      that takes one argument (`device_idx`) and should return the value that is
      going to be created on devices[device_idx].
    devices: a list of device strings to create `PerReplica` values on.

  Returns:
    A `PerReplica` object.
  """
  values = []
  for device_idx, device in enumerate(devices):
    if callable(value):
      v = value(device_idx)
    elif isinstance(value, list):
      v = value[device_idx]
    else:
      v = value
    if isinstance(v, IndexedSlicesValue):
      with ops.device(device):
        values.append(
            IndexedSlices(
                values=array_ops.identity(v.values),
                indices=array_ops.identity(v.indices),
                dense_shape=array_ops.identity(v.dense_shape)))
    else:
      with ops.device(device):
        values.append(array_ops.identity(v))
  return value_lib.PerReplica(values)


def enable_collective_ops():
  """Enable collectives in the current process."""
  cluster_resolver = cluster_resolver_lib.TFConfigClusterResolver()
  context.context().configure_collective_ops(
      collective_leader="'/job:worker/replica:0/task:0'")
  config_proto = config_pb2.ConfigProto()
  config_proto.experimental.collective_group_leader = (
      "/job:worker/replica:0/task:0")
  server_def = tensorflow_server_pb2.ServerDef(
      cluster=cluster_resolver.cluster_spec().as_cluster_def(),
      default_session_config=config_proto,
      job_name=cluster_resolver.task_type,
      task_index=cluster_resolver.task_id,
      protocol=cluster_resolver.rpc_layer)
  context.context().enable_collective_ops(server_def)
  # Recover default flag values.
  CollectiveReplicaLauncher._prefer_unique_instance_key = True
  CollectiveReplicaLauncher._prefer_ordering_token = False


class MultiProcessPoolRunner():

  def __init__(self, num_processes):
    cluster_spec_dict = multi_worker_test_base.create_cluster_spec(
        num_workers=num_processes)
    self.runner = multi_process_runner.MultiProcessPoolRunner(cluster_spec_dict)


# Global MultiProcessPoolRunners that can be shared by test cases to avoid
# expensive initialization cost of TensorFlow in new processes.
#
# Note that they have to be globals and can't be owned by test classes because
# usually fn usually captures the test class instance, and test class
# instance can't be pickled if it has mpr as a member (it is not allowed to
# pickle Process objects).
# TODO(crccw): Use `num_workers` combination once it is ready.
global_mpr_2p = MultiProcessPoolRunner(num_processes=2)
global_mpr_1p = MultiProcessPoolRunner(num_processes=1)


def get_global_mpr(num_processes):
  if num_processes == 1:
    return global_mpr_1p.runner
  elif num_processes == 2:
    return global_mpr_2p.runner
  else:
    raise ValueError("get_global_mpr: num_processes must be 1 or 2, got %d" %
                     num_processes)


class CollectiveOpsTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super().setUp()
    # Enabling collectives can be done in "setUpClass", but requires using
    # different collective_keys in different tests as collectives are reused
    # across tests. Always resetting collective ops before each test offers
    # better test isolation.
    global_mpr_1p.runner.run(enable_collective_ops)
    global_mpr_2p.runner.run(enable_collective_ops)

  def make_collective(self, num_processes, gpu_per_process):
    """Returns collectives and other info to be used in tests.

    Args:
      num_processes: an integer indicating the number of processes that
        participate in the collective.
      gpu_per_process: number of GPUs (0 if no GPUs) used by each process.

    Returns:
     A tuple of (collective, devices, pid) where collective is a instance
     of `CollectiveAllReduce`, devices are a list of local devices (str)
     attached to the current process, and pid is the id of this process among
     all participant processes.
    """

    cluster_resolver = cluster_resolver_lib.TFConfigClusterResolver()
    devices = [
        "/job:worker/replica:0/task:%d/device:CPU:0" % cluster_resolver.task_id
    ]
    if gpu_per_process > 0:
      devices = [
          "/job:worker/replica:0/task:%d/device:GPU:%d" %
          (cluster_resolver.task_id, i) for i in range(gpu_per_process)
      ]
    group_size = num_processes * len(devices)
    collective = cross_device_ops_lib.CollectiveAllReduce(
        devices=devices,
        group_size=group_size,
        options=collective_util.Options())
    return collective, devices, cluster_resolver.task_id

  def as_list(self, value):
    """An utility to convert a `Mirrored`, `Tensor` or `IndexedSlices` to a list.

    The reason it exists is to provide a uniformed view of returned value of
    "reduce" calls, especially across tf.function boundaries. Returning
    `Mirrored` from a tf.function will only evaluate the primary value, which
    makes collective ops of non-primary device being pruned, and will eventually
    cause hanging.

    Args:
      value: the value to convert, can be one of `Mirrored`, `Tensor` and
        `IndexedSlices`.

    Returns:
      A list of `Tensor` or `IndexedSlices`.
    """
    if isinstance(value, ops.Tensor):
      return [value]
    elif isinstance(value, IndexedSlices):
      return [value]
    elif isinstance(value, value_lib.Mirrored):
      return value.values
    else:
      raise ValueError("unwrap: unsupported input type: %s" % type(value))

  RunOptions = collections.namedtuple(  # pylint: disable=invalid-name
      "RunOptions",
      [
          "mode",  # A list of str from ["eager", "func_graph"]
          "num_processes",
          "gpus_per_process",
          "reduce_op",
          "communication_options",
          "prefer_unique_instance_key",
      ])
  RunOptions.__new__.__defaults__ = (["eager",
                                      "func_graph"], 2, 0, ReduceOp.SUM,
                                     collective_util.Options(), True)

  def reduce_and_verify(self, inputs, expect, options):
    """Reduce the given `inputs` and verify the output matches `expect`.

    Args:
      inputs: a list of `Tensor` or `IndexedSlices`, where i-th value will be
        fed to i-th replica.
      expect: a `Tensor` or `IndexedSlices`. This should be the expected value
        for one replica.
      options: a `RunOpotions` instance.
    """

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          options.prefer_unique_instance_key)
      collective, devices, pid = self.make_collective(options.num_processes,
                                                      options.gpus_per_process)

      def reduce_fn():
        value_fn = lambda device_idx: inputs[pid * len(devices) + device_idx]
        per_replica_value = make_per_replica_value(value_fn, devices)
        reduced_values = collective.reduce(options.reduce_op, per_replica_value,
                                           per_replica_value,
                                           options.communication_options)
        if options.gpus_per_process > 1:
          self.assertIsInstance(reduced_values, value_lib.Mirrored)
        reduced_values = self.as_list(reduced_values)
        self.assertAllEqual(devices, [v.device for v in reduced_values])
        return [ops.convert_to_tensor(v) for v in reduced_values]

      per_replica_expect = [ops.convert_to_tensor(expect)] * len(devices)

      if "eager" in options.mode:
        got = reduce_fn()
        self.assertAllClose(got, per_replica_expect)

      if "func_graph" in options.mode:
        got = def_function.function(reduce_fn)()
        self.assertAllClose(got, per_replica_expect)

    get_global_mpr(options.num_processes).run(replica_fn)

  def batch_reduce_and_verify(self, inputs, expect, options):
    """Batch reduce the given `inputs` and verify the output matches `expect`.

    Args:
      inputs: a 2-level nested list of `Tensor` or `IndexedSlices`, where i-th
        value will be fed to i-th replica.
      expect: a list of `Tensor` or `IndexedSlices`. This should be the expected
        value for one replica.
      options: a `RunOpotions` instance.
    """

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          options.prefer_unique_instance_key)
      collective, devices, pid = self.make_collective(options.num_processes,
                                                      options.gpus_per_process)

      def batch_reduce_fn():
        batch_size = len(inputs[0])
        value_dst_pairs = []
        for i in range(batch_size):

          def value_fn(device_idx, idx=i):
            return inputs[pid * len(devices) + device_idx][idx]

          per_replica_value = make_per_replica_value(value_fn, devices)
          value_dst_pairs.append((per_replica_value, per_replica_value))
        reduced_values = collective.batch_reduce(options.reduce_op,
                                                 value_dst_pairs,
                                                 options.communication_options)
        if options.gpus_per_process > 1:
          for v in reduced_values:
            self.assertIsInstance(v, value_lib.Mirrored)
        reduced_values = [self.as_list(v) for v in reduced_values]
        for v in reduced_values:
          self.assertAllEqual(devices, [t.device for t in v])
        return nest.map_structure(ops.convert_to_tensor, reduced_values)

      per_replica_expect = nest.map_structure(
          lambda x: [ops.convert_to_tensor(x)] * len(devices), expect)

      if "eager" in options.mode:
        got = batch_reduce_fn()
        self.assertAllClose(got, per_replica_expect)

      if "func_graph" in options.mode:
        got = def_function.function(batch_reduce_fn)()
        self.assertAllClose(got, per_replica_expect)

    get_global_mpr(options.num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          reduce_op=[ReduceOp.SUM, ReduceOp.MEAN],
          prefer_unique_instance_key=[True, False]))
  def testReduceDense(self, num_processes, required_gpus, implementation,
                      reduce_op, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")
    options = self.RunOptions(
        num_processes=num_processes,
        gpus_per_process=required_gpus,
        reduce_op=reduce_op,
        communication_options=collective_util.Options(
            implementation=implementation),
        prefer_unique_instance_key=prefer_unique_instance_key)
    group_size = options.num_processes * (options.gpus_per_process or 1)

    inputs_data = [1.0, 2.0, 3.0, 4.0]
    inputs = inputs_data[0:group_size]

    if group_size == 1:
      expect = 1.0
    if group_size == 2:
      expect = 3.0 if reduce_op == ReduceOp.SUM else 1.5
    elif group_size == 4:
      expect = 10.0 if reduce_op == ReduceOp.SUM else 2.5

    self.reduce_and_verify(inputs, expect, options)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          # TODO(b/166682130): add MEAN reduce once the bug is fixed.
          reduce_op=ReduceOp.SUM,
          prefer_unique_instance_key=[True, False]))
  def testReduceSparse(self, num_processes, required_gpus, implementation,
                       reduce_op, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")
    options = self.RunOptions(
        mode=["func_graph"],  # Sparse reduce is not supported in eager.
        num_processes=num_processes,
        gpus_per_process=required_gpus,
        reduce_op=reduce_op,
        communication_options=collective_util.Options(
            implementation=implementation),
        prefer_unique_instance_key=prefer_unique_instance_key)
    group_size = options.num_processes * (options.gpus_per_process or 1)

    inputs_data = [
        IndexedSlicesValue(
            values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[3.], [4.]], indices=[1, 2], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[5.], [6.]], indices=[7, 8], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[7.], [8.]], indices=[3, 2], dense_shape=[10, 1]),
    ]
    inputs = inputs_data[0:group_size]

    if group_size == 1:
      expect = IndexedSlices(
          values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1])
    elif group_size == 2:
      expect = IndexedSlices(
          values=[[1.], [2.], [3.], [4.]],
          indices=[0, 1, 1, 2],
          dense_shape=[10, 1])
    elif group_size == 4:
      expect = IndexedSlices(
          values=[[1.], [2.], [3.], [4.], [5.], [6.], [7.], [8.]],
          indices=[0, 1, 1, 2, 7, 8, 3, 2],
          dense_shape=[10, 1])

    self.reduce_and_verify(inputs, expect, options)

  @combinations.generate(
      combinations.combine(prefer_unique_instance_key=[True, False]))
  def testReduceSparseVariableLength(self, prefer_unique_instance_key):
    # One device per process, 2 processes, 2 replicas in total.
    inputs = [
        IndexedSlicesValue(values=[[1.]], indices=[0], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[2.], [3.], [4.]], indices=[0, 1, 2], dense_shape=[10, 1]),
    ]
    expect = IndexedSlices(
        values=[[1.], [2.], [3.], [4.]],
        indices=[0, 0, 1, 2],
        dense_shape=[10, 1])
    self.reduce_and_verify(
        inputs,
        expect,
        self.RunOptions(
            mode=["func_graph"],  # Sparse reduce is not supported in eager.
            num_processes=2,
            reduce_op=ReduceOp.SUM,
            prefer_unique_instance_key=prefer_unique_instance_key))

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          reduce_op=[ReduceOp.SUM, ReduceOp.MEAN],
          prefer_unique_instance_key=[True, False]))
  def testBatchReduceDense(self, num_processes, required_gpus, implementation,
                           reduce_op, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    options = self.RunOptions(
        num_processes=num_processes,
        gpus_per_process=required_gpus,
        reduce_op=reduce_op,
        communication_options=collective_util.Options(
            implementation=implementation),
        prefer_unique_instance_key=prefer_unique_instance_key)
    group_size = options.num_processes * (options.gpus_per_process or 1)

    inputs_data = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]
    inputs = inputs_data[0:group_size]

    if group_size == 1:
      expect = [1.0, 2.0]
    if group_size == 2:
      expect = [4.0, 6.0] if reduce_op == ReduceOp.SUM else [2.0, 3.0]
    elif group_size == 4:
      expect = [16.0, 20.0] if reduce_op == ReduceOp.SUM else [4.0, 5.0]

    self.batch_reduce_and_verify(inputs, expect, options)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          # TODO(b/166682130): add MEAN reduce once the bug is fixed.
          reduce_op=ReduceOp.SUM,
          prefer_unique_instance_key=[True, False]))
  def testBatchReduceSparse(self, num_processes, required_gpus, implementation,
                            reduce_op, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    options = self.RunOptions(
        mode=["func_graph"],  # Sparse reduce is not supported in eager.
        num_processes=num_processes,
        gpus_per_process=required_gpus,
        reduce_op=reduce_op,
        communication_options=collective_util.Options(
            implementation=implementation),
        prefer_unique_instance_key=prefer_unique_instance_key)
    group_size = options.num_processes * (options.gpus_per_process or 1)

    inputs_data = ([
        IndexedSlicesValue(
            values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[3.], [4.]], indices=[1, 2], dense_shape=[5, 1])
    ], [
        IndexedSlicesValue(
            values=[[5.], [6.]], indices=[1, 2], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[7.], [8.]], indices=[0, 1], dense_shape=[5, 1])
    ], [
        IndexedSlicesValue(
            values=[[9.], [10.]], indices=[3, 4], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[11.], [12.]], indices=[3, 4], dense_shape=[5, 1])
    ], [
        IndexedSlicesValue(
            values=[[13.], [14.]], indices=[8, 9], dense_shape=[10, 1]),
        IndexedSlicesValue(
            values=[[15.], [16.]], indices=[3, 4], dense_shape=[5, 1])
    ])
    inputs = inputs_data[0:group_size]

    if group_size == 1:
      expect = [
          IndexedSlices(
              values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]),
          IndexedSlices(
              values=[[3.], [4.]], indices=[1, 2], dense_shape=[5, 1])
      ]
    if group_size == 2:
      expect = [
          IndexedSlices(
              values=[[1.], [2.], [5.], [6.]],
              indices=[0, 1, 1, 2],
              dense_shape=[10, 1]),
          IndexedSlices(
              values=[[3.], [4.], [7.], [8.]],
              indices=[1, 2, 0, 1],
              dense_shape=[5, 1])
      ]
    elif group_size == 4:
      expect = [
          IndexedSlices(
              values=[[1.], [2.], [5.], [6.], [9.], [10.], [13.], [14.]],
              indices=[0, 1, 1, 2, 3, 4, 8, 9],
              dense_shape=[10, 1]),
          IndexedSlices(
              values=[[3.], [4.], [7.], [8.], [11.], [12.], [15.], [16.]],
              indices=[1, 2, 0, 1, 3, 4, 3, 4],
              dense_shape=[5, 2])
      ]
    self.batch_reduce_and_verify(inputs, expect, options)

  def testBatchReduceMixedDenseAndSparse(self):

    options = self.RunOptions(
        num_processes=2,
        gpus_per_process=0,
        reduce_op=ReduceOp.SUM,
        mode=["func_graph"])

    inputs_data = [
        [
            1.0, 2.0,
            IndexedSlicesValue(
                values=[[1.], [2.]], indices=[0, 1], dense_shape=[10, 1]),
            IndexedSlicesValue(
                values=[[3.], [4.]], indices=[1, 2], dense_shape=[5, 1])
        ],
        [
            3.0, 4.0,
            IndexedSlicesValue(
                values=[[5.], [6.]], indices=[1, 2], dense_shape=[10, 1]),
            IndexedSlicesValue(
                values=[[7.], [8.]], indices=[0, 1], dense_shape=[5, 1])
        ],
    ]

    expect = [
        4.0, 6.0,
        IndexedSlices(
            values=[[1.], [2.], [5.], [6.]],
            indices=[0, 1, 1, 2],
            dense_shape=[10, 1]),
        IndexedSlices(
            values=[[3.], [4.], [7.], [8.]],
            indices=[1, 2, 0, 1],
            dense_shape=[5, 1])
    ]

    self.batch_reduce_and_verify(inputs_data, expect, options)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          reduce_op=[ReduceOp.SUM, ReduceOp.MEAN],
      ))
  def testAllReduceDense(self, num_processes, required_gpus, implementation,
                         reduce_op):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)
      group_size = num_processes * (required_gpus or 1)

      @def_function.function
      def collective_all_reduce():
        results = []
        for replica_id, device in enumerate(devices):
          with ops.device(device):
            value = constant_op.constant(1.0)
            results.append(
                collective._all_reduce(reduce_op, value, replica_id, options))
        return results

      got = collective_all_reduce()
      if reduce_op == ReduceOp.SUM:
        expect = [1.0 * group_size] * len(devices)
      elif reduce_op == ReduceOp.MEAN:
        expect = [1.0] * len(devices)
      self.assertAllClose(got, expect)

      @def_function.function
      def collective_batch_all_reduce():
        results = []
        for replica_id, device in enumerate(devices):
          with ops.device(device):
            value = (constant_op.constant(1.0), constant_op.constant(2.0))
            results.append(
                collective._all_reduce(reduce_op, value, replica_id, options))
        return results

      got = collective_batch_all_reduce()
      if reduce_op == ReduceOp.SUM:
        expect = [(1.0 * group_size, 2.0 * group_size)] * len(devices)
      elif reduce_op == ReduceOp.MEAN:
        expect = [(1.0, 2.0)] * len(devices)
      self.assertAllClose(got, expect)

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          reduce_op=[ReduceOp.SUM, ReduceOp.MEAN],
      ))
  def testAllReduceSparse(self, num_processes, required_gpus, implementation,
                          reduce_op):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)
      group_size = num_processes * (required_gpus or 1)

      @def_function.function
      def collective_all_reduce():
        results = []
        for replica_id, device in enumerate(devices):
          with ops.device(device):
            value = IndexedSlices(
                values=array_ops.identity([[1.]]),
                indices=array_ops.identity([0]),
                dense_shape=array_ops.identity([5, 1]))
            results.append(
                collective._all_reduce(reduce_op, value, replica_id, options))
        return results

      got = collective_all_reduce()
      if reduce_op == ReduceOp.SUM:
        expect = [IndexedSlices([[1. * group_size]], [0], [5, 1])
                 ] * len(devices)
      elif reduce_op == ReduceOp.MEAN:
        expect = [IndexedSlices([[1.]], [0], [5, 1])] * len(devices)
      self.assertAllClose(
          nest.map_structure(ops.convert_to_tensor, got),
          nest.map_structure(ops.convert_to_tensor, expect))

      @def_function.function
      def collective_batch_all_reduce():
        results = []
        for replica_id, device in enumerate(devices):
          with ops.device(device):
            value = (IndexedSlices(
                array_ops.identity([[1.]]), array_ops.identity([0]),
                array_ops.identity([5, 1])),
                     IndexedSlices(
                         array_ops.identity([[3.]]), array_ops.identity([2]),
                         array_ops.identity([5, 1])))
            results.append(
                collective._all_reduce(reduce_op, value, replica_id, options))
        return results

      got = collective_batch_all_reduce()
      if reduce_op == ReduceOp.SUM:
        expect = [(IndexedSlices([[1. * group_size]], [0], [5, 1]),
                   IndexedSlices([[3. * group_size]], [2], [5, 1]))
                 ] * len(devices)
      elif reduce_op == ReduceOp.MEAN:
        expect = [(IndexedSlices([[1.]], [0], [5, 1]),
                   IndexedSlices([[3.]], [2], [5, 1]))] * len(devices)
      self.assertAllClose(
          nest.map_structure(ops.convert_to_tensor, got),
          nest.map_structure(ops.convert_to_tensor, expect))

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=2,
          required_gpus=0,
          implementation=CommunicationImplementation.AUTO,
          reduce_op=ReduceOp.SUM))
  def testAllReduceMixedDenseAndSparse(self, num_processes, required_gpus,
                                       implementation, reduce_op):

    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)
      group_size = num_processes * (required_gpus or 1)

      @def_function.function
      def collective_batch_all_reduce():
        results = []
        for replica_id, device in enumerate(devices):
          with ops.device(device):
            value = (IndexedSlices(
                array_ops.identity([[1.]]), array_ops.identity([0]),
                array_ops.identity([5, 1])), array_ops.identity(1.0),
                     IndexedSlices(
                         array_ops.identity([[3.]]), array_ops.identity([2]),
                         array_ops.identity([5, 1])), array_ops.identity(2.0))
            results.append(
                collective._all_reduce(reduce_op, value, replica_id, options))
        return results

      got = collective_batch_all_reduce()
      expect = [
          (IndexedSlices([[1. * group_size]], [0], [5, 1]), 1.0 * group_size,
           IndexedSlices([[3. * group_size]], [2], [5, 1]), 2.0 * group_size)
      ] * len(devices)
      self.assertAllClose(
          nest.map_structure(ops.convert_to_tensor, got),
          nest.map_structure(ops.convert_to_tensor, expect))

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          axis=[0, 1, 2],
          func_mode=["eager", "func_graph"],
          implementation=[
              CommunicationImplementation.AUTO,
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testAllGatherSameShape(self, num_processes, required_gpus, implementation,
                             func_mode, axis, prefer_unique_instance_key):

    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)
      value = constant_op.constant([[[1, 2], [1, 2]]], dtype=dtypes.float32)

      def gather_fn():
        per_replica_value = make_per_replica_value(value, devices)
        gathered_values = collective._gather(
            per_replica_value, per_replica_value, axis=axis, options=options)
        gathered_values = self.as_list(gathered_values)
        # Skip checking devices in eager. In eager the device attribute doesn't
        # reflect the actual device of the tensor.
        if not context.executing_eagerly():
          self.assertAllEqual(devices, [v.device for v in gathered_values])
        return [ops.convert_to_tensor(v) for v in gathered_values]

      group_size = num_processes * (required_gpus or 1)
      expect = array_ops.concat([value] * group_size, axis=axis)
      per_replica_expect = [ops.convert_to_tensor(expect)] * len(devices)

      if func_mode == "eager":
        result = gather_fn()
        self.assertAllClose(result, per_replica_expect)

      if func_mode == "func_graph":
        result = def_function.function(gather_fn)()
        self.assertAllClose(result, per_replica_expect)

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=[1, 2],
          required_gpus=[0, 1, 2],
          implementation=[CommunicationImplementation.RING]))
  def testCollectiveV2ControlFlow(self, num_processes, required_gpus,
                                  implementation):

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = True
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)
      value = make_per_replica_value(constant_op.constant([1.]), devices)

      @def_function.function
      def reduce_fn():

        def cond_body():
          reduced = collective.reduce(reduce_util.ReduceOp.SUM, value, value,
                                      options)
          return math_ops.add_n(self.as_list(reduced)) / len(devices)

        return cond.cond(
            array_ops.identity(False), cond_body, cond_body)

      num_replicas = num_processes * len(devices)
      self.assertAllEqual(reduce_fn(), [1. * num_replicas])

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=1,
          required_gpus=2,
          implementation=[
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testMultiThreadedCollectiveLaunchNoInterleave(self, num_processes,
                                                    required_gpus,
                                                    implementation,
                                                    prefer_unique_instance_key):

    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)

      # We would like to simulate the following sequence:
      #   thread-0  device0                 device1
      #   thread-1          device0 device1
      # If the kernel launch sequence is as-is the program will deadlock since
      # NCCL requires the launch order to be same on each device.
      v0 = make_per_replica_value(1.0, devices)
      v1 = make_per_replica_value(2.0, devices)

      # Add a delay to collective_ops.all_reduce according to the input tensors
      # index in `sequence.`
      sequence = [v0.values[0], v1.values[0], v1.values[1], v0.values[1]]
      all_reduce = collective_ops.all_reduce

      def delayed_all_reduce(input_tensor, *args, **kwargs):
        for idx, v in enumerate(sequence):
          if input_tensor is v:
            time.sleep(idx)
            break
        return all_reduce(input_tensor, *args, **kwargs)

      with test.mock.patch.object(collective_ops, "all_reduce",
                                  delayed_all_reduce):
        # We only use NCCL for batch reduce with two or more values, so we use
        # two values here.

        def thread_fn():
          reduced = collective.batch_reduce(reduce_util.ReduceOp.SUM,
                                            [(v0, v0), (v0, v0)], options)
          self.assertAllEqual(reduced[0].values, [2.0, 2.0])
          self.assertAllEqual(reduced[1].values, [2.0, 2.0])

        t = threading.Thread(target=thread_fn)
        t.start()
        reduced = collective.batch_reduce(reduce_util.ReduceOp.SUM, [(v1, v1),
                                                                     (v1, v1)],
                                          options)
        self.assertAllEqual(reduced[0].values, [4.0, 4.0])
        self.assertAllEqual(reduced[1].values, [4.0, 4.0])
        t.join()

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=1,
          required_gpus=2,
          implementation=[
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testInputsAreFunctionArgs(self, num_processes, required_gpus,
                                implementation, prefer_unique_instance_key):

    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(implementation=implementation)

      @def_function.function
      def reduce_fn(v):
        # Function inputs don't have device placement.
        self.assertEqual(v.values[0].device, "")
        self.assertEqual(v.values[1].device, "")
        # We only use NCCL for batch reduce with two or more values, so we use
        # two values here.
        reduced = collective.batch_reduce(reduce_util.ReduceOp.SUM, [(v, v),
                                                                     (v, v)],
                                          options)
        self.assertEqual(reduced[0].values[0].device, devices[0])
        self.assertEqual(reduced[0].values[1].device, devices[1])
        self.assertEqual(reduced[1].values[0].device, devices[0])
        self.assertEqual(reduced[1].values[1].device, devices[1])
        # Returning Mirrored only evaluates the primary value, which causes
        # hanging,
        return [reduced[0].values, reduced[1].values]

      v = make_per_replica_value(1.0, devices)
      reduced = reduce_fn(v)
      self.assertAllClose(reduced, [[2.0, 2.0], [2.0, 2.0]])

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=2,
          required_gpus=[0, 1],
          implementation=[
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testTimeoutReduceDense(self, num_processes, implementation, required_gpus,
                             prefer_unique_instance_key):

    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, task_id = self.make_collective(
          num_processes, required_gpus)
      if task_id != 0:
        return

      v = make_per_replica_value(1.0, devices)
      options = collective_util.Options(
          timeout_seconds=1., implementation=implementation)

      @def_function.function
      def reduce_dense():
        return collective.reduce(reduce_util.ReduceOp.SUM, v, v, options)

      # The collective should time out because we only launch it on worker-0,
      # while there're three workers in total.
      with self.assertRaises(errors.DeadlineExceededError):
        reduce_dense()

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=2,
          required_gpus=[0, 1],
          implementation=[
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testTimeoutBatchReduceDense(self, num_processes, implementation,
                                  required_gpus, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, task_id = self.make_collective(
          num_processes, required_gpus)
      if task_id != 0:
        return

      v = make_per_replica_value(1.0, devices)
      options = collective_util.Options(
          timeout_seconds=1., implementation=implementation)

      @def_function.function
      def batch_reduce_dense():
        return collective.batch_reduce(reduce_util.ReduceOp.SUM,
                                       [(v, v), (v, v)], options)

      # The collective should time out because we only launch it on worker-0,
      # while there're two workers in total.
      with self.assertRaises(errors.DeadlineExceededError):
        batch_reduce_dense()

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=2,
          required_gpus=[0, 1],
          implementation=[
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testTimeoutReduceSparse(self, num_processes, implementation,
                              required_gpus, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, task_id = self.make_collective(
          num_processes, required_gpus)
      if task_id != 0:
        return

      v = make_per_replica_value(
          IndexedSlicesValue(
              values=[[4., 6.]], indices=[1], dense_shape=[5, 2]), devices)
      options = collective_util.Options(
          timeout_seconds=1., implementation=implementation)

      @def_function.function
      def reduce_sparse():
        return collective.reduce(reduce_util.ReduceOp.SUM, v, v, options)

      # The collective should time out because we only launch it on worker-0,
      # while there're two workers in total.
      with self.assertRaises(errors.DeadlineExceededError):
        reduce_sparse()

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(
      combinations.combine(
          num_processes=2,
          required_gpus=[0, 1],
          implementation=[
              CommunicationImplementation.RING,
              CommunicationImplementation.NCCL,
          ],
          prefer_unique_instance_key=[True, False]))
  def testTimeoutBatchReduceSparse(self, num_processes, required_gpus,
                                   implementation, prefer_unique_instance_key):
    if (required_gpus == 0 and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip CPU + NCCL combination")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.NCCL):
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")
    if (num_processes != required_gpus and
        implementation == CommunicationImplementation.AUTO):
      self.skipTest("Skip potential NCCL combination (AUTO) with mismatched "
                    "process and GPU count. NCCL requires physical GPUs for "
                    "every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = (
          prefer_unique_instance_key)
      collective, devices, task_id = self.make_collective(
          num_processes, required_gpus)
      if task_id != 0:
        return

      v = make_per_replica_value(
          IndexedSlicesValue(
              values=[[4., 6.]], indices=[1], dense_shape=[5, 2]), devices)
      options = collective_util.Options(
          timeout_seconds=1., implementation=implementation)

      @def_function.function
      def batch_reduce_sparse():
        return collective.batch_reduce(reduce_util.ReduceOp.SUM,
                                       [(v, v), (v, v)], options)

      # The collective should time out because we only launch it on worker-0,
      # while there're two workers in total.
      with self.assertRaises(errors.DeadlineExceededError):
        batch_reduce_sparse()

    get_global_mpr(num_processes).run(replica_fn)

  @combinations.generate(combinations.combine(num_processes=1, required_gpus=2))
  def testNcclOrdering(self, num_processes, required_gpus):

    if num_processes != required_gpus:
      self.skipTest("Skip NCCL combination with mismatched process and GPU "
                    "count. NCCL requires physical GPUs for every process.")

    def replica_fn():
      CollectiveReplicaLauncher._prefer_unique_instance_key = True
      CollectiveReplicaLauncher._prefer_ordering_token = True
      collective, devices, _ = self.make_collective(num_processes,
                                                    required_gpus)
      options = collective_util.Options(
          implementation=CommunicationImplementation.NCCL)

      v_dense = make_per_replica_value([1.0, 1.0], devices)
      v_sparse = make_per_replica_value([
          IndexedSlicesValue([[4., 6.], [5., 6.]], [1, 3], [5, 2]),
          IndexedSlicesValue([[4., 6.], [5., 6.]], [1, 3], [5, 2]),
      ], devices)

      @def_function.function
      def nested_dense():
        collective.reduce(reduce_util.ReduceOp.SUM, v_dense, v_dense, options)

      @def_function.function
      def nested_sparse():
        collective.reduce(reduce_util.ReduceOp.SUM, v_sparse, v_sparse, options)

      # All collectives, function calls, if clause and while loops should be
      # chained by control dependencies, so that the execution order is
      # deterministic.
      @def_function.function
      def f():
        # pylint: disable=pointless-statement
        collective.reduce(reduce_util.ReduceOp.SUM, v_sparse, v_sparse, options)
        # reducing dense value.
        collective.reduce(reduce_util.ReduceOp.SUM, v_dense, v_dense, options)
        # reducing sparse value.
        collective.reduce(reduce_util.ReduceOp.SUM, v_sparse, v_sparse, options)
        # reduce dense value in nested tf.function.
        nested_dense()
        # reduce sparse value in nested tf.function.
        nested_sparse()
        # reduce dense value in tf.cond.
        if array_ops.identity(1.0) > array_ops.identity(2.0):
          collective.reduce(reduce_util.ReduceOp.SUM, v_dense, v_dense, options)
        else:
          v_dense
        # reduce sparse value in tf.cond.
        if array_ops.identity(1.0) > array_ops.identity(2.0):
          v_sparse
        else:
          collective.reduce(reduce_util.ReduceOp.SUM, v_sparse, v_sparse,
                            options)
        # reduce dense value in tf.while_loop.
        i = array_ops.identity(1)
        while i < 3:
          collective.reduce(reduce_util.ReduceOp.SUM, v_dense, v_dense, options)
          i += 1
        # reduce sparse value in tf.while_loop.
        i = array_ops.identity(1)
        while i < 3:
          collective.reduce(reduce_util.ReduceOp.SUM, v_sparse, v_sparse,
                            options)
          i += 1
        # reducing dense and sparse value again.
        collective.reduce(reduce_util.ReduceOp.SUM, v_dense, v_dense, options)
        collective.reduce(reduce_util.ReduceOp.SUM, v_sparse, v_sparse, options)
        # pylint: enable=pointless-statement

      graph = f.get_concrete_function().graph
      should_be_ordered = set([
          "CollectiveReduceV2", "CollectiveGatherV2", "If", "While",
          "StatefulPartitionedCall"
      ])
      nodes_by_device = {}
      for op in graph.get_operations():
        if op.type in should_be_ordered:
          if op.device not in nodes_by_device:
            nodes_by_device[op.device] = []
          nodes_by_device[op.device].append(op)
      order = test_util.topological_sort_operations(graph.get_operations())
      for device in devices:
        device = device_util.canonicalize(device)
        # Those function ops don't have device annotations, but they contain
        # collectives for both devices so we always include them.
        operations = nodes_by_device[device] + nodes_by_device[""]
        # Verify that we get all types of nodes we want.
        self.assertEqual(set(op.type for op in operations), should_be_ordered)
        test_util.assert_sequential_execution(order, operations)

    get_global_mpr(num_processes).run(replica_fn)


@_REGISTER_DECORATOR(CollectiveOpsTest)
def _save_test_case(pickler, obj):
  def reconstruct(*args, **kwargs):
    del args, kwargs
    return CollectiveOpsTest()

  return pickler.save_reduce(reconstruct, (), obj=obj)


if __name__ == "__main__":
  # Set default inter op thread pool size to one to ensure we don't exhaust the
  # thread pool with the additional executors to run collectives in eager.
  os.environ["TF_NUM_INTEROP_THREADS"] = "1"
  # TODO(b/172304955): figure why logical devices doesn't work.
  test_util.main(config_logical_devices=False)