code: def testAllGatherGradientNest(self, strategy, pure_eager):
      if pure_eager:
          self.skipTest(
              "`tf.gradients` is not supported with eager execution "
              "without using tf.functions."
          )
  
      def all_gather_fn(value):
          axis = 1
          ctx = distribute_lib.get_replica_context()
          return ctx.all_gather(array_ops.identity(value), axis)
  
      gradient_comp = sum(range(1, strategy.num_replicas_in_sync + 1))
      gradient = [[gradient_comp], [gradient_comp]]
      grads_for_all_replicas = [
          gradient
      ] * _get_num_replicas_per_client(strategy)
  
      @def_function.function
      def step(c):
          x = constant_op.constant([[3.0], [5.0]])
          y = constant_op.constant([[2.0], [4.0]])
          mid = all_gather_fn([x, y])
          y = mid * c
          return gradients_impl.gradients_v2(y, [x])[0]
  
      def value_fn(ctx):
          x = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
          return array_ops.constant([x[ctx.replica_id_in_sync_group]])
  
      per_replica_value = (
          strategy.experimental_distribute_values_from_function(
              value_fn
          )
      )
      result = strategy.experimental_local_results(
          strategy.run(step, args=(per_replica_value,))
      )
  
      self.assertAllEqual(grads_for_all_replicas, 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 29 for calling the GatherTest obj's underlying member method code with example usage: obj.testAllGatherGradientNest(strategy, pure_eager) and returns: ctx, gradients_impl, array_ops