code: def test_log_prob_ratio_with_no_axes(self):
      # The motivation here is to enable writing generic code that works inside
      # and outside a pmap.
      def lp_ratio_grad(x0, x1, axis_name):
          dist0 = sharded.Sharded(
              sample_dist_lib.Sample(normal.Normal(0.0, 1.0), x0.shape),
              shard_axis_name=axis_name,
          )
          dist1 = sharded.Sharded(
              sample_dist_lib.Sample(normal.Normal(0.0, 2.0), x0.shape),
              shard_axis_name=axis_name,
          )
  
          return gradient.value_and_gradient(
              lambda x0, x1: log_prob_ratio.log_prob_ratio(
                  dist0, x0, dist1, x1
              ),
              (x0, x1),
          )
  
      def true_lp_ratio_grad(x0, x1):
          dist0 = sample_dist_lib.Sample(
              normal.Normal(0.0, 1.0), x0.shape
          )
          dist1 = sample_dist_lib.Sample(
              normal.Normal(0.0, 2.0), x0.shape
          )
  
          return gradient.value_and_gradient(
              lambda x0, x1: log_prob_ratio.log_prob_ratio(
                  dist0, x0, dist1, x1
              ),
              (x0, x1),
          )
  
      x0 = tf.range(1, 1 + test_lib.NUM_DEVICES, dtype=tf.float32)
      x1 = tf.range(5, 5 + test_lib.NUM_DEVICES, dtype=tf.float32)
      lp_ratio, (g0, g1) = lp_ratio_grad(x0, x1, axis_name=[])
      sharded_x0 = self.shard_values(x0)
      sharded_x1 = self.shard_values(x1)
      sharded_lp_ratio, (
          sharded_g0,
          sharded_g1,
      ) = self.per_replica_to_tensor(
          self.strategy_run(
              lambda x0, x1: lp_ratio_grad(x0, x1, self.axis_name),
              (sharded_x0, sharded_x1),
          )
      )
      true_lp_ratio, (true_g0, true_g1) = true_lp_ratio_grad(x0, x1)
  
      self.assertAllClose(true_lp_ratio, lp_ratio)
      self.assertAllClose(true_g0, g0)
      self.assertAllClose(true_g1, g1)
      self.assertAllClose(true_lp_ratio, sharded_lp_ratio[0])
      self.assertAllClose(true_g0, sharded_g0)
      self.assertAllClose(true_g1, sharded_g1)
  
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 9 for calling the ShardTest obj's underlying member method code with example usage: obj.test_log_prob_ratio_with_no_axes() and returns: gradient, gradient