code: def test_gradient_is_correctly_reduced_with_multiple_axes(self):
      if not JAX_MODE:
          self.skipTest("Multiple shard axes not supported in TF.")
  
      other_axis_name = self.axis_name + "_other"
  
      def run(x, y, z):
          def log_prob_parts(value):
              x, y, z = value
              return [
                  normal.Normal(0.0, 1.0).log_prob(x),
                  normal.Normal(0.0, 1.0).log_prob(y),
                  normal.Normal(x + y, 1.0).log_prob(z),
              ]
  
          def log_prob(x, y, z):
              sharded_log_prob_parts = (
                  distribute_lib.make_sharded_log_prob_parts(
                      log_prob_parts,
                      [
                          self.axis_name,
                          other_axis_name,
                          [self.axis_name, other_axis_name],
                      ],
                  )
              )
              parts = sharded_log_prob_parts([x, y, z])
              return tf.add_n(parts)
  
          return gradient.value_and_gradient(log_prob, (x, y, z))
  
      seed = random.PRNGKey(0)
      x_seed, y_seed, z_seed = samplers.split_seed(seed, n=3)
      x = normal.Normal(0.0, 1).sample(seed=x_seed, sample_shape=2)
      y = normal.Normal(0.0, 1.0).sample(seed=y_seed, sample_shape=2)
      z = normal.Normal(0.0, 1.0).sample(
          seed=z_seed, sample_shape=[2, 2]
      )
  
      def outer_run(x, y, z):
          return self.strategy_run(
              run,
              (x, y, z),
              in_axes=(None, 0, 0),
              axis_name=other_axis_name,
          )
  
      out_values, out_grads = self.strategy_run(
          outer_run, (x, y, z), in_axes=(0, None, 0)
      )
  
      def true_log_prob(x, y, z):
          return (
              tf.reduce_sum(normal.Normal(0.0, 1.0).log_prob(x))
              + tf.reduce_sum(normal.Normal(0.0, 1.0).log_prob(y))
              + tf.reduce_sum(
                  normal.Normal(x[:, None] + y[None], 1.0).log_prob(z)
              )
          )
  
      true_values, true_grads = self.evaluate(
          gradient.value_and_gradient(true_log_prob, (x, y, z))
      )
  
      self.assertAllClose(
          out_values,
          tf.ones([2, 2]) * true_values,
          rtol=1e-6,
          atol=1e-6,
      )
      self.assertAllEqualNested(
          out_grads[0], tf.ones([2, 2]) * true_grads[0][:, None]
      )
      self.assertAllEqualNested(
          out_grads[1], tf.ones([2, 2]) * true_grads[1][None]
      )
      self.assertAllEqualNested(
          out_grads[2], tf.ones([2, 2]) * true_grads[2]
      )
  
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 16 for calling the LogProbPartsTest obj's underlying member method code with example usage: obj.test_gradient_is_correctly_reduced_with_multiple_axes() and returns: gradient, self, unknown, list, tf