code: def test_jd_log_prob_ratio(self, dist_fn):
      dist = dist_fn(self.axis_name)
  
      def run(key):
          sample = dist.sample(seed=key)
          log_prob = dist.log_prob(sample)
          return sample, log_prob
  
      keys = samplers.split_seed(self.key, 2)
      samples = []
      unmapped_samples = []
      log_probs = []
      true_log_probs = []
  
      for key in keys:
          sample, log_prob = self.per_replica_to_tensor(
              self.strategy_run(run, (key,), in_axes=None)
          )
  
          if isinstance(dist, jd.JointDistributionNamed):
              # N.B. the global RV 'w' gets replicated, so we grab any single
              # replica's result.
              w, x, data = sample["w"][0], sample["x"], sample["data"]
          else:
              w, x, data = sample[0][0], sample[1], sample[2]
  
          true_log_prob = true_log_prob_fn(w, x, data)
  
          samples.append(sample)
          unmapped_samples.append((w, x, data))
          log_probs.append(log_prob[0])
          true_log_probs.append(true_log_prob)
  
      def true_diff(x, y):
          return true_log_prob_fn(*x) - true_log_prob_fn(*y)
  
      def run_diff(x, y):
          def _lpr(x, y):
              return log_prob_ratio.log_prob_ratio(dist, x, dist, y)
  
          return gradient.value_and_gradient(_lpr, [x, y])
  
      dist_lp_diff, dist_lp_diff_grad = self.per_replica_to_tensor(
          self.strategy_run(
              run_diff,
              tuple(tf.nest.map_structure(self.shard_values, samples)),
          )
      )
  
      true_lp_diff, true_lp_diff_grad = gradient.value_and_gradient(
          true_diff, unmapped_samples
      )
  
      if isinstance(dist, jd.JointDistributionNamed):
          dist_lp_diff_grad[0] = (
              dist_lp_diff_grad[0]["w"][0],
              dist_lp_diff_grad[0]["x"],
              dist_lp_diff_grad[0]["data"],
          )
          dist_lp_diff_grad[1] = (
              dist_lp_diff_grad[1]["w"][0],
              dist_lp_diff_grad[1]["x"],
              dist_lp_diff_grad[1]["data"],
          )
      else:
          true_lp_diff_grad[0] = list(true_lp_diff_grad[0])
          true_lp_diff_grad[1] = list(true_lp_diff_grad[1])
          dist_lp_diff_grad[0] = list(dist_lp_diff_grad[0])
          dist_lp_diff_grad[0][0] = dist_lp_diff_grad[0][0][0]
          dist_lp_diff_grad[1] = list(dist_lp_diff_grad[1])
          dist_lp_diff_grad[1][0] = dist_lp_diff_grad[1][0][0]
  
      lp_diff = log_probs[0] - log_probs[1]
  
      self.assertAllClose(
          true_lp_diff, lp_diff, rtol=7e-6
      )  # relaxed tol for fp32 in JAX
      self.assertAllClose(true_lp_diff, dist_lp_diff[0])
      self.assertAllCloseNested(true_lp_diff_grad, dist_lp_diff_grad)
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import scale
  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 lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.distribute import joint_distribution
  from tensorflow_probability.python.experimental.distribute import sharded
  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 5 for calling the JointDistributionTest obj's underlying member method code with example usage: obj.test_jd_log_prob_ratio(dist_fn) and returns: sample,  log_prob, unknown, gradient, log_prob_ratio