code: def test_all_reduce(self, enable_mixed_precision_reduce):
      if enable_mixed_precision_reduce:
          os.environ["DTENSOR_ENABLE_MIXED_PRECISION_REDUCE"] = ""
          # Override group size since we are testing on smaller mesh.
          os.environ["DTENSOR_REDUCE_IN_BFLOAT16_MAX_GROUP_SIZE"] = "4"
      else:
          if "DTENSOR_ENABLE_MIXED_PRECISION_REDUCE" in os.environ:
              del os.environ["DTENSOR_ENABLE_MIXED_PRECISION_REDUCE"]
  
      @polymorphic_function.function
      def _compute_reduction(inp):
          return math_ops.reduce_sum(inp, axis=[2])
  
      input_tensor = stateless_random_ops.stateless_random_uniform(
          shape=(8, 8, 8, 64),
          seed=self.stateless_random_seed,
          minval=-5.0,
          maxval=5.0,
          dtype=dtypes.bfloat16,
      )
      expected = _compute_reduction(input_tensor)
  
      # Compute reduction on 8x1, since dim 2 is unsharded AllReduce will not be
      # needed.
      mesh_8x1 = self._create_mesh((8, 1), "TPU")
      input_8x1 = numpy_util.pack_numpy(
          input_tensor,
          Layout(
              [_MESH_DIM_X, UNSHARDED, UNSHARDED, UNSHARDED], mesh_8x1
          ),
      )
      result_8x1 = _compute_reduction(input_8x1)
      result_8x1_np = numpy_util.to_numpy(result_8x1)
  
      # Compute reduction on 1x8, AllReduce will be needed since dim 2 is sharded.
      mesh_1x8 = self._create_mesh((1, 8), "TPU")
      input_1x8 = numpy_util.pack_numpy(
          input_tensor,
          Layout(
              [_MESH_DIM_X, UNSHARDED, _MESH_DIM_Y, UNSHARDED], mesh_1x8
          ),
      )
      result_1x8 = _compute_reduction(input_1x8)
      result_1x8_np = numpy_util.to_numpy(result_1x8)
  
      self.assertEqual(result_8x1.dtype, dtypes.bfloat16)
      self.assertEqual(result_1x8.dtype, dtypes.bfloat16)
  
      # Mixed precision does not apply since AllReduce was not used, result will
      # always be close to the expected value.
      self.assertAllClose(result_8x1_np, expected, atol=1e-5, rtol=1e-5)
  
      # AllReduce was needed, so result will be more accurate if mixed precision
      # is enabled.
      if enable_mixed_precision_reduce:
          self.assertAllClose(
              result_1x8_np, expected, atol=1e-5, rtol=1e-5
          )
      else:
          self.assertNotAllClose(
              result_1x8_np, expected, atol=1e-5, rtol=1e-5
          )
  
file_path: tensorflow/tensorflow/dtensor/python/tests/numerics_test.py
imports: import os
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import accelerator_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 3 for calling the NumericTest obj's underlying member method code with example usage: obj.test_all_reduce(enable_mixed_precision_reduce) and returns: math_ops