code: def testWhileWithScopedAllocator(self):
      group_size = 2
      group_key = 1
      instance_key0 = 1
      instance_key1 = 2
  
      config = config_pb2.ConfigProto(device_count={"CPU": group_size})
      rewrite_options = config.graph_options.rewrite_options
      rewrite_options.scoped_allocator_optimization = (
          rewriter_config_pb2.RewriterConfig.ON
      )
      del rewrite_options.scoped_allocator_opts.enable_op[:]
      rewrite_options.scoped_allocator_opts.enable_op.append(
          "CollectiveReduce"
      )
  
      # Tests that execute collectives need to be enclosed in graph or tf.function
      with ops.Graph().as_default():
          with self.session(config=config) as sess:
              run_ops = []
              for i in range(group_size):
                  with ops.device("CPU:%d" % i):
                      constant = constant_op.constant(0.0)
                      cond = lambda i: math_ops.less(i, 10.0)
                      body = lambda i: math_ops.add(i, 1.0)
                      input0 = while_loop.while_loop(
                          cond, body, [constant]
                      )
                      input1 = math_ops.add(constant, 5)
                      colred0 = collective_ops.all_reduce(
                          input0,
                          group_size,
                          group_key,
                          instance_key0,
                          "Add",
                          "Id",
                      )
                      colred1 = collective_ops.all_reduce(
                          input1,
                          group_size,
                          group_key,
                          instance_key1,
                          "Add",
                          "Id",
                      )
                      run_ops.append(math_ops.add_n([colred0, colred1]))
              results = sess.run(run_ops)
          self.assertEqual(results, [30.0, 30.0])
  
imports: import time
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.core.protobuf import rewriter_config_pb2
  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 dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import kernels
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import collective_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_v1
  from tensorflow.python.ops import variables
  from tensorflow.python.ops import while_loop
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
label_desc: Use image node_id 13 for calling the CollectiveOpTest obj's underlying member method code with example usage: obj.testWhileWithScopedAllocator() without return types