docstr: all_gather(..., axis=0,...) a nest of DistributedValues.
code: def testAllGatherNest1D0Axis(self, strategy, pure_eager):
      """all_gather(..., axis=0,...) a nest of DistributedValues."""
      single_value = constant_op.constant([1, 2, 3])
      axis = 0
  
      def run():
          value_identity = array_ops.identity(single_value)
          ctx = distribute_lib.get_replica_context()
          return ctx.all_gather(
              [value_identity, value_identity], axis=axis
          )
  
      if not pure_eager:
          run = def_function.function(run)
  
      all_value = [
          single_value for _ in range(strategy.num_replicas_in_sync)
      ]
      expect = array_ops.concat(all_value, axis=axis)
      expected_per_replica = [expect] * _get_num_replicas_per_client(
          strategy
      )
  
      result = strategy.run(run)
      for gathered_result in result:
          self.assertAllEqual(
              expected_per_replica,
              strategy.experimental_local_results(gathered_result),
          )
  
imports: from absl.testing import parameterized
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import central_storage_strategy
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import mirrored_strategy
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute.collective_all_reduce_strategy import CollectiveAllReduceStrategy
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import indexed_slices
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gradients_impl
  from tensorflow.python.platform import test
  from tensorflow.python.util import nest
label_desc: Use image node_id 24 for calling the GatherTest obj's underlying member method code with example usage: obj.testAllGatherNest1D0Axis(strategy, pure_eager) and returns: ctx