code: def testSubgroupAllReduceCombiner(self):
      self.skipForDeviceType(
          ["TPU"],
          "This test requires 8 TPU cores.",
          unless_device_count_equals_to=8,
      )
  
      # Create and use an eight-device mesh just for this test.
      global_ids = test_util.create_device_ids_array((4, 2))
      local_ids = np.ravel(global_ids).tolist()
      mesh_dict = {
          device: layout_lib.Mesh(
              _MESH_DIMS,
              global_ids,
              local_ids,
              test_util.create_device_list((4, 2), device),
          )
          for device in ("CPU", "GPU", "TPU")
      }
  
      mesh = self.configTestMesh(mesh_dict)
      fully_sharded_layout_2d = Layout(_MESH_DIMS, mesh)
  
      @polymorphic_function.function
      def func(a, b):
          a = math_ops.reduce_sum(a, axis=[0])
          b = math_ops.reduce_sum(b, axis=[0])
          # Do something with the results before adding them, to make sure the MLIR
          # pass can handle dependent ops sandwiched between two all-reduce ops.
          return gen_math_ops.square(a) + gen_math_ops.square(b)
  
      row = constant_op.constant(
          np.array([[1.0, 2.0]]), dtype=dtypes.float32
      )
      a = array_ops.repeat(row, repeats=[8], axis=0)
      b = gen_array_ops.reverse_v2(a, axis=[1])
      expected_result = func(a, b)
  
      a = api.relayout(a, fully_sharded_layout_2d)
      b = api.relayout(b, fully_sharded_layout_2d)
      dtensor_result = func(a, b)
  
      self.assertDTensorEqual(
          expected_result, Layout([_MESH_DIM_Y], mesh), dtensor_result
      )
  
imports: import os
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import dtensor_device
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 3 for calling the CollectiveTestWithCustomMesh obj's underlying member method code with example usage: obj.testSubgroupAllReduceCombiner() and returns: unknown