code: def test_reduce_with_collective_grads_matches_without_collectives(
      self, reduce_op, distributed_op, is_supported, skip_on_eager=False
  ):
      if not is_supported:
          self.skipTest("Gradient of operation not supported.")
      if skip_on_eager and (tf.executing_eagerly() and TF_MODE):
          self.skipTest("Not supported in Eager.")
      x = tf.reshape(
          tf.range(test_lib.NUM_DEVICES * 6.0) / 5.0,
          [test_lib.NUM_DEVICES, 3, 2],
      )
  
      def compute_dist_grads(x):
          return gradient.value_and_gradient(
              lambda x: distributed_op(  # pylint: disable=g-long-lambda
                  x, axis=[0, 1], named_axis=self.axis_name
              ),
              [x],
          )[1][0]
  
      def distributed_run(x):
          return self.per_replica_to_tensor(
              self.strategy_run(
                  compute_dist_grads, (self.shard_values(x),)
              )
          )
  
      reduce_grads = gradient.value_and_gradient(
          lambda x: reduce_op(x, axis=None), [x]
      )[1][0]
      dist_grads = distributed_run(x)
      self.assertAllClose(reduce_grads, dist_grads)
  
imports: import functools
  import itertools
  from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 4 for calling the CollectiveTest obj's underlying member method code with example usage: obj.test_reduce_with_collective_grads_matches_without_collectives(reduce_op, distributed_op, is_supported, skip_on_eager) and returns: gradient, self