code: def test_jd_has_correct_sample_path_gradients_with_partial_values(
      self,
  ):
      def run(seed):
          @jd.JointDistributionCoroutine
          def model():
              yield Root(normal.Normal(0.0, 1.0, name="x"))
              yield normal.Normal(0.0, 1.0, name="y")
              yield sharded.Sharded(
                  normal.Normal(1.0, 1.0), self.axis_name, name="z"
              )
  
          sample = model.sample(seed=seed)
  
          def lp_fn1(x, y, z):
              return model.log_prob((x, y, z))
  
          def lp_fn2(x, z):
              return model.log_prob(
                  model.sample(value=(x, None, z), seed=seed)
              )
  
          lp_and_grad1 = gradient.value_and_gradient(lp_fn1, [*sample])
          (lp2, grad2) = gradient.value_and_gradient(
              lp_fn2, [sample.x, sample.z]
          )
          return lp_and_grad1, (lp2, grad2)
  
      (lp1, grad1), (lp2, grad2) = self.per_replica_to_tensor(
          self.strategy_run(run, (self.key,), in_axes=None)
      )
      grad2 = [grad2[0], None, grad2[1]]
      for i in range(test_lib.NUM_DEVICES):
          for j in range(3):
              self.assertAllClose(lp1[i], lp2[i])
              if grad2[j] is not None:
                  self.assertAllClose(grad1[j][i], grad2[j][i])
  
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 7 for calling the JointDistributionTest obj's underlying member method code with example usage: obj.test_jd_has_correct_sample_path_gradients_with_partial_values() and returns: lp_and_grad1,  lp2,  grad2, model, model