docstr: Different at non-`axis`-th dimension : [1, 1], [1, 2], 0th -> raise error.
code: def testGatherRaiseDiffShapeAtNonAxis(self, strategy, pure_eager):
      """Different at non-`axis`-th dimension : [1, 1], [1, 2], 0th -> raise error."""
      if (
          isinstance(strategy, CollectiveAllReduceStrategy)
          and _get_num_replicas_per_client(strategy) > 1
      ):
          self.skipTest("b/167331966")
  
      if strategy.num_replicas_in_sync <= 1:
          self.skipTest("Test for more than 1 replica only.")
  
      def value_fn(ctx):
          return constant_op.constant(
              1, shape=(1, ctx.replica_id_in_sync_group + 1)
          )
  
      distributed_values = (
          strategy.experimental_distribute_values_from_function(
              value_fn
          )
      )
      axis = 0
  
      def run():
          return strategy.gather(distributed_values, axis=axis)
  
      if not pure_eager:
          run = def_function.function(run)
  
      if isinstance(strategy, CollectiveAllReduceStrategy):
          with self.assertRaisesRegex(
              errors.InvalidArgumentError, r"Shape mismatch"
          ):
              run()
      elif isinstance(
          strategy,
          (
              mirrored_strategy.MirroredStrategy,
              central_storage_strategy.CentralStorageStrategy,
          ),
      ):
          with self.assertRaisesRegex(
              (errors.InvalidArgumentError, ValueError),
              r"Dimension \d in both shapes must be equal",
          ):
              run()
  
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 10 for calling the GatherTest obj's underlying member method code with example usage: obj.testGatherRaiseDiffShapeAtNonAxis(strategy, pure_eager) and returns: constant_op, strategy