code: 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)
  
imports: 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
  from tensorflow.python.distribute import collective_util
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import cross_device_ops
  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
  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.framework import tensor
  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
label_desc: Use image node_id 15 for calling the CollectiveOpsTest obj's underlying member method code with example usage: obj.testAllGatherSameShape(num_processes, required_gpus, implementation, func_mode, axis, prefer_unique_instance_key) and returns: unknown