code: def test_batch_and_event_shape_with_plate(self, jd_class):
      models = {}
  
      def coroutine_model():
          g = yield lognormal.LogNormal(0.0, 1.0)
          df = yield exponential.Exponential(1.0)
          loc = yield sample_lib.Sample(normal.Normal(0, g), 20)
          yield student_t.StudentT(tf.expand_dims(df, -1), loc, 1)
  
      models[
          jdab.JointDistributionCoroutineAutoBatched
      ] = coroutine_model
  
      models[jdab.JointDistributionSequentialAutoBatched] = [
          lognormal.LogNormal(0.0, 1.0),
          exponential.Exponential(1.0),
          lambda _, g: sample_lib.Sample(normal.Normal(0, g), 20),
          lambda loc, df: student_t.StudentT(
              tf.expand_dims(df, -1), loc, 1
          ),
      ]
  
      models[
          jdab.JointDistributionNamedAutoBatched
      ] = collections.OrderedDict(
          (
              ("g", lognormal.LogNormal(0.0, 1.0)),
              ("df", exponential.Exponential(1.0)),
              (
                  "loc",
                  lambda g: sample_lib.Sample(normal.Normal(0, g), 20),
              ),
              (
                  "x",
                  lambda loc, df: student_t.StudentT(
                      tf.expand_dims(df, -1), loc, 1
                  ),
              ),
          )
      )
  
      joint = jd_class(models[jd_class], validate_args=True)
  
      # Properties `event_shape` and `batch_shape` should be defined
      # even before any sampling calls have occurred.
      self.assertAllEqualNested(
          joint._model_flatten(joint.event_shape), [[], [], [20], [20]]
      )
      self.assertAllEqual(joint.batch_shape, [])
  
      is_scalar = joint._model_flatten(joint.is_scalar_event())
      self.assertAllEqual(is_scalar[0], True)
      self.assertAllEqual(is_scalar[1], True)
      self.assertAllEqual(is_scalar[2], False)
      self.assertAllEqual(is_scalar[3], False)
  
      event_shape = joint._model_flatten(joint.event_shape_tensor())
      self.assertAllEqual(event_shape[0], [])
      self.assertAllEqual(event_shape[1], [])
      self.assertAllEqual(event_shape[2], [20])
      self.assertAllEqual(event_shape[3], [20])
  
      self.assertEqual(joint.is_scalar_batch(), True)
  
      batch_shape = joint.batch_shape_tensor()
      self.assertAllEqual(batch_shape, [])
  
      self.assertIsInstance(joint, tf.__internal__.CompositeTensor)
      flat = tf.nest.flatten(joint, expand_composites=True)
      unflat = tf.nest.pack_sequence_as(
          joint, flat, expand_composites=True
      )
      self.assertIsInstance(unflat, jd_class)
      self.assertIs(type(joint.model), type(unflat.model))
  
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 1 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_batch_and_event_shape_with_plate(jd_class) without return types