code: def test_distributed_reduce_works_as_normal_with_int_axes(
      self, reduce_op, distributed_op, skip_on_eager=False
  ):
      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 make_run(axis):
          return lambda x: distributed_op(x, axis=axis)
  
      for axes in [0, 1, None, [0, 1]]:
          reduce_out = reduce_op(x, axis=axes)
          dist_out = make_run(axes)(x)
          self.assertAllEqual(reduce_out, dist_out)
  
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 2 for calling the CollectiveTest obj's underlying member method code with example usage: obj.test_distributed_reduce_works_as_normal_with_int_axes(reduce_op, distributed_op, skip_on_eager) and returns: function