code: def test_nested_shard_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, data):
          def log_prob_parts(value):
              x, data = value
              return [
                  normal.Normal(0.0, 1.0).log_prob(x),
                  normal.Normal(x, 1.0).log_prob(data),
              ]
  
          def log_prob(x, data):
              sharded_log_prob_parts = (
                  distribute_lib.make_sharded_log_prob_parts(
                      log_prob_parts,
                      [None, {self.axis_name, other_axis_name}],
                  )
              )
              parts = sharded_log_prob_parts([x, data])
              return tf.add_n(parts)
  
          return gradient.value_and_gradient(log_prob, (x, data))
  
      x = tf.constant(1.0)
      data = 2 * tf.ones([2, 2])
  
      def outer_run(x, data):
          return self.strategy_run(
              run,
              (x, data),
              in_axes=(None, 0),
              axis_name=other_axis_name,
          )
  
      out_values, out_grads = self.strategy_run(
          outer_run, (x, data), in_axes=(None, 0)
      )
  
      def true_log_prob(x, data):
          return normal.Normal(0.0, 1.0).log_prob(x) + tf.reduce_sum(
              normal.Normal(x, 1.0).log_prob(data)
          )
  
      true_values, true_grads = self.evaluate(
          gradient.value_and_gradient(true_log_prob, (x, data))
      )
  
      self.assertAllEqualNested(
          out_values, tf.ones([2, 2]) * true_values
      )
      self.assertAllEqualNested(
          out_grads[0], tf.ones([2, 2]) * true_grads[0]
      )
      self.assertAllEqualNested(
          out_grads[1], tf.ones([2, 2]) * true_grads[1]
      )
  
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 15 for calling the LogProbPartsTest obj's underlying member method code with example usage: obj.test_nested_shard_axes() and returns: gradient, self, unknown, list, tf