code: def test_log_prob_ratio_sample(self):
      dist1 = sample_dist_lib.Sample(
          normal.Normal(0.0, 4.0), test_lib.NUM_DEVICES
      )
      dist2 = sample_dist_lib.Sample(
          normal.Normal(0.0, 2.0), test_lib.NUM_DEVICES
      )
  
      x0 = tf.zeros([test_lib.NUM_DEVICES])
      x1 = tf.ones([test_lib.NUM_DEVICES])
  
      sharded_x0 = self.shard_values(x0)
      sharded_x1 = self.shard_values(x1)
  
      true_diff, (true_g1, true_g2) = gradient.value_and_gradient(
          lambda x0, x1: log_prob_ratio.log_prob_ratio(
              dist1, x0, dist2, x1
          ),
          (x0, x1),
      )
  
      @tf.function
      def test(x0, x1):
          dist_sharded1 = sharded.Sharded(
              sample_dist_lib.Sample(normal.Normal(0.0, 4.0), 1),
              shard_axis_name=self.axis_name,
          )
          dist_sharded2 = sharded.Sharded(
              sample_dist_lib.Sample(normal.Normal(0.0, 2.0), 1),
              shard_axis_name=self.axis_name,
          )
          return (
              gradient.value_and_gradient(
                  lambda x0, x1: log_prob_ratio.log_prob_ratio(  # pylint: disable=g-long-lambda
                      dist_sharded1, x0, dist_sharded2, x1
                  ),
                  (x0, x1),
              ),
              dist_sharded1.log_prob(x0) - dist_sharded2.log_prob(x1),
          )
  
      (diff1, (g1, g2)), diff2 = self.per_replica_to_tensor(
          self.strategy_run(test, (sharded_x0, sharded_x1))
      )
  
      self.assertAllClose(true_diff, diff1[0])
      self.assertAllClose(true_diff, diff2[0])
      self.assertAllClose(true_g1, g1)
      self.assertAllClose(true_g2, g2)
  
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 10 for calling the ShardTest obj's underlying member method code with example usage: obj.test_log_prob_ratio_sample() and returns: gradient,  unknown