code: def test_jd(self, dist_fn):
      dist = dist_fn(self.axis_name)
  
      def run(key):
          sample = dist.sample(seed=key)
          # The identity is to prevent reparameterization gradients from kicking in.
          log_prob, (log_prob_grads,) = gradient.value_and_gradient(
              dist.log_prob,
              (tf.nest.map_structure(tf.identity, sample),),
          )
          return sample, log_prob, log_prob_grads
  
      sample, log_prob, log_prob_grads = self.strategy_run(
          run, (self.key,), in_axes=None
      )
      sample, log_prob, log_prob_grads = self.per_replica_to_tensor(
          (sample, log_prob, log_prob_grads)
      )
  
      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"]
          log_prob_grads = (
              log_prob_grads["w"][0],
              log_prob_grads["x"],
              log_prob_grads["data"],
          )
      else:
          w, x, data = sample[0][0], sample[1], sample[2]
          log_prob_grads = (
              log_prob_grads[0][0],
              log_prob_grads[1],
              log_prob_grads[2],
          )
  
      true_log_prob, true_log_prob_grads = gradient.value_and_gradient(
          true_log_prob_fn, (w, x, data)
      )
  
      self.assertAllClose(
          self.evaluate(log_prob),
          self.evaluate(tf.ones(test_lib.NUM_DEVICES) * true_log_prob),
      )
      self.assertAllCloseNested(
          self.evaluate(log_prob_grads),
          self.evaluate(true_log_prob_grads),
      )
  
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 4 for calling the JointDistributionTest obj's underlying member method code with example usage: obj.test_jd(dist_fn) and returns: sample,  log_prob,  log_prob_grads