code: def test_gather_dtensor_value(self):
      strategy = mwms.MultiWorkerMirroredStrategy()
      stride = self.num_client * self.num_local_devices
  
      def value_fn(value_context):
          start = value_context.replica_id_in_sync_group * stride
          return array_ops.reshape(
              math_ops.range(start=start, limit=start + stride),
              shape=(1, stride),
          )
  
      distribute_result = (
          strategy.experimental_distribute_values_from_function(
              value_fn
          )
      )
      # distribute_result is a DTensorDistributedValue.
      # The shape of the global tensor is [stride, stride],
      # and each worker gets [stride/2, stride].
      result = strategy.gather(distribute_result, axis=0)
  
      start = stride * self.num_local_devices * self.client_id
      end = start + stride * self.num_local_devices
      self.assertEqual(result.shape, [self.num_local_devices, stride])
      self.assertAllClose(
          result,
          array_ops.reshape(
              math_ops.range(start=start, limit=end),
              shape=(self.num_local_devices, -1),
          ),
      )
  
      result = strategy.gather(distribute_result, axis=1)
      self.assertEqual(
          result.shape, [1, self.num_local_devices * stride]
      )
      self.assertAllClose(
          result,
          array_ops.reshape(
              math_ops.range(start=start, limit=end), shape=(1, -1)
          ),
      )
  
imports: import json
  import os
  from absl import flags
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import multi_client_test_util
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import reduce_util
  from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
  from tensorflow.python.distribute.experimental import dtensor_util
  from tensorflow.python.distribute.experimental import multi_worker_mirrored_strategy
  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 tensor_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 19 for calling the MultiWorkerMirroredStrategyTest obj's underlying member method code with example usage: obj.test_gather_dtensor_value() and returns: array_ops