code: def test_jd_has_correct_sample_path_gradients(self):
      def log_prob_fn(x_loc):
          @jdc.JointDistributionCoroutine
          def surrogate():
              x = yield Root(normal.Normal(x_loc, 1.0))
              y = yield normal.Normal(x, 1.0)
              yield sample_dist_lib.Sample(
                  normal.Normal(x + y, 1.0), test_lib.NUM_DEVICES
              )
  
          @jdc.JointDistributionCoroutine
          def model():
              yield Root(normal.Normal(1.0, 1.0))
              yield Root(normal.Normal(1.0, 1.0))
              yield sample_dist_lib.Sample(
                  normal.Normal(1.0, 1.0), test_lib.NUM_DEVICES
              )
  
          return tf.reduce_mean(
              model.log_prob(
                  surrogate.sample(sample_shape=1e5, seed=self.key)
              )
          )
  
      true_log_prob, true_log_prob_grad = gradient.value_and_gradient(
          log_prob_fn, 0.0
      )
  
      def run(seed):
          def sharded_log_prob_fn(x_loc):
              @jd.JointDistributionCoroutine
              def surrogate():
                  x = yield Root(normal.Normal(x_loc, 1.0))
                  y = yield normal.Normal(x, 1.0)
                  yield sharded.Sharded(
                      normal.Normal(x + y, 1.0), self.axis_name
                  )
  
              @jd.JointDistributionCoroutine
              def model():
                  yield Root(normal.Normal(1.0, 1.0))
                  yield Root(normal.Normal(1.0, 1.0))
                  yield sharded.Sharded(
                      normal.Normal(1.0, 1.0), self.axis_name
                  )
  
              return tf.reduce_mean(
                  model.log_prob(
                      surrogate.sample(sample_shape=1e5, seed=seed)
                  )
              )
  
          (
              sharded_log_prob,
              sharded_log_prob_grad,
          ) = gradient.value_and_gradient(sharded_log_prob_fn, 0.0)
          return sharded_log_prob, sharded_log_prob_grad
  
      (
          sharded_log_prob,
          sharded_log_prob_grad,
      ) = self.per_replica_to_tensor(
          self.strategy_run(run, (self.key,), in_axes=None)
      )
      for i in range(test_lib.NUM_DEVICES):
          self.assertAllClose(
              sharded_log_prob[i], true_log_prob, atol=0.025
          )
          self.assertAllClose(
              sharded_log_prob_grad[i], true_log_prob_grad, atol=2e-2
          )
  
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 6 for calling the JointDistributionTest obj's underlying member method code with example usage: obj.test_jd_has_correct_sample_path_gradients() and returns: tf, sharded_log_prob,  sharded_log_prob_grad, tf