code: def testCompiledAllReduce(self):
      self._setup_context()
  
      def all_reduce_sum(v):
          return collective_ops.all_reduce_v2(
              t=v,
              group_size=2,
              group_key=1,
              instance_key=1,
              merge_op="Add",
              final_op="Id",
          )
  
      strategy = mirrored_strategy.MirroredStrategy(["GPU:0", "GPU:1"])
  
      @def_function.function(jit_compile=True)
      def f():
          return while_loop.while_loop(
              lambda i, _: i < 5,
              lambda i, t: (i + 1, all_reduce_sum(t)),
              (array_ops.zeros([]), constant_op.constant(1.0)),
          )
  
      @def_function.function
      def run():
          return strategy.run(f)
  
      _, reduce = strategy.experimental_local_results(run())[0]
      self.assertEqual(reduce.numpy(), 32.0)
  
file_path: tensorflow-upstream/tensorflow/python/ops/compiled_collective_ops_gpu_test.py
imports: import os
  from tensorflow.python.distribute import mirrored_strategy
  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 ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import collective_ops
  from tensorflow.python.ops import while_loop
  from tensorflow.python.platform import test
label_desc: Use image node_id 3 for calling the CompiledCollectiveOpGPUTest obj's underlying member method code with example usage: obj.testCompiledAllReduce() and returns: collective_ops, while_loop, strategy