code: 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)
  
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 18 for calling the CollectiveOpsTest obj's underlying member method code with example usage: obj.testInputsAreFunctionArgs(num_processes, required_gpus, implementation, prefer_unique_instance_key) and returns: list