docstr: Different at non-`axis`-th dimension : [2, 1], [1, 1], all_gather(...axis=1...) -> raise error.
code: def testAllGatherRaiseDiffShapeAtNonAxis(self, strategy, pure_eager):
      """Different at non-`axis`-th dimension : [2, 1], [1, 1], all_gather(...axis=1...) -> raise error."""
      if _is_tpu_strategy(strategy):
          self.skipTest(
              "TODO(b/169108777): raise a clear error message in xla."
          )
  
      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)
          )
  
      per_replica_value = (
          strategy.experimental_distribute_values_from_function(
              value_fn
          )
      )
  
      def run(value):
          value_identity = array_ops.identity(value)
          ctx = distribute_lib.get_replica_context()
          return ctx.all_gather(value_identity, axis=0)
  
      if not pure_eager:
          run = def_function.function(run)
  
      if isinstance(strategy, CollectiveAllReduceStrategy):
          with self.assertRaisesRegex(
              errors.InvalidArgumentError, r"Shape mismatch"
          ):
              strategy.run(run, args=(per_replica_value,))
      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",
          ):
              strategy.run(run, args=(per_replica_value,))
  
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 25 for calling the GatherTest obj's underlying member method code with example usage: obj.testAllGatherRaiseDiffShapeAtNonAxis(strategy, pure_eager) and returns: constant_op, ctx