code: def testAllGatherDiffValueTPU(self, strategy, pure_eager):
      # Test for TPU only since it can't be tested via testAllGatherDiffShape*
      if not _is_tpu_strategy(strategy):
          self.skipTest(
              "Test for TPU only. For other strategies case already"
              " covered in other tests"
          )
  
      data = [[1], [2], [3], [4], [5], [6], [7], [8]]
  
      axis = 0
      dataset = dataset_ops.DatasetV2.from_tensor_slices(data).batch(8)
      input_iterator = iter(
          strategy.experimental_distribute_dataset(dataset)
      )
  
      @def_function.function
      def replica_fn(per_replica_value):
          ctx = distribute_lib.get_replica_context()
          return ctx.all_gather(
              array_ops.identity(per_replica_value), axis=axis
          )
  
      result = strategy.experimental_local_results(
          strategy.run(replica_fn, args=(next(input_iterator),))
      )
  
      expected_result = [data] * _get_num_replicas_per_client(strategy)
      self.assertAllClose(expected_result, 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 20 for calling the GatherTest obj's underlying member method code with example usage: obj.testAllGatherDiffValueTPU(strategy, pure_eager) and returns: ctx