code: def test_log_prob_with_no_axes(self):
      # The motivation here is to enable writing generic code that works inside
      # and outside a pmap.
      def lp_grad(x, axis_name):
          return gradient.value_and_gradient(
              sharded.Sharded(
                  sample_dist_lib.Sample(
                      normal.Normal(0.0, 1.0), x.shape
                  ),
                  shard_axis_name=axis_name,
              ).log_prob,
              (x,),
          )
  
      def true_lp_grad(x):
          return gradient.value_and_gradient(
              sample_dist_lib.Sample(
                  normal.Normal(0.0, 1.0), x.shape
              ).log_prob,
              (x,),
          )
  
      x = tf.range(1, 1 + test_lib.NUM_DEVICES, dtype=tf.float32)
      lp, g = lp_grad(x, axis_name=[])
      sharded_x = self.shard_values(x)
      sharded_lp, sharded_g = self.per_replica_to_tensor(
          self.strategy_run(
              lambda x: lp_grad(x, self.axis_name), (sharded_x,)
          )
      )
      true_lp, true_g = true_lp_grad(x)
  
      self.assertAllClose(true_lp, lp)
      self.assertAllCloseNested(true_g, g)
      self.assertAllClose(true_lp, sharded_lp[0])
      self.assertAllCloseNested(true_g, sharded_g)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.distribute import joint_distribution
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.experimental.distributions import increment_log_prob
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 8 for calling the ShardTest obj's underlying member method code with example usage: obj.test_log_prob_with_no_axes() and returns: gradient, gradient