code: def test_nested_joint_distributions(self):
      batch_shape = [2, 3]
  
      def inner_fn():
          xy = yield jdab.JointDistributionNamedAutoBatched(
              {
                  "x": normal.Normal(
                      loc=tf.zeros(batch_shape),
                      scale=tf.ones(batch_shape),
                      name="x",
                  ),
                  "y": lambda x: poisson.Poisson(log_rate=x, name="y"),
              },
              batch_ndims=2,
              name="xy",
          )
          _ = yield normal.Normal(loc=0.0, scale=xy["y"], name="z")
  
      joint = jdab.JointDistributionSequentialAutoBatched(
          [
              jdab.JointDistributionCoroutineAutoBatched(
                  inner_fn, batch_ndims=1, name="a"
              )
          ]
      )
      z = joint.sample(seed=test_util.test_seed())
  
      # Batch and event shape.
      self.assertAllEqual(joint.batch_shape, [])
      self.assertAllEqualNested(
          tf.nest.map_structure(lambda x: tf.TensorShape(x.shape), z),
          joint.event_shape,
      )
  
      # Sample shape.
      z2 = self.evaluate(joint.sample(5, seed=test_util.test_seed()))
      lp2 = joint.log_prob(z2)
      self.assertAllEqual(lp2.shape, [5])
  
      z3 = joint.sample(value=z2, seed=test_util.test_seed())
      self.assertAllCloseNested(z2, z3)
  
imports: import collections
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 21 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_nested_joint_distributions() without return types