code: def test_log_prob_and_prob_with_plate(self, jd_class):
      models = {}
  
      def coroutine_model():
          a = yield bernoulli.Bernoulli(probs=0.5, dtype=tf.float32)
          b = yield sample_lib.Sample(  # pylint: disable=g-long-lambda
              bernoulli.Bernoulli(
                  probs=0.25 + 0.5 * a, dtype=tf.float32
              ),
              2,
          )
          yield normal.Normal(loc=a, scale=1.0 + b)
  
      models[
          jdab.JointDistributionCoroutineAutoBatched
      ] = coroutine_model
      models[jdab.JointDistributionSequentialAutoBatched] = [
          bernoulli.Bernoulli(probs=0.5, dtype=tf.float32),
          lambda a: sample_lib.Sample(  # pylint: disable=g-long-lambda
              bernoulli.Bernoulli(
                  probs=0.25 + 0.5 * a, dtype=tf.float32
              ),
              2,
          ),
          lambda b, a: normal.Normal(loc=a, scale=1.0 + b),
      ]
      models[
          jdab.JointDistributionNamedAutoBatched
      ] = collections.OrderedDict(
          (
              ("a", bernoulli.Bernoulli(probs=0.5, dtype=tf.float32)),
              (
                  "b",
                  lambda a: sample_lib.Sample(  # pylint: disable=g-long-lambda
                      bernoulli.Bernoulli(
                          probs=0.25 + 0.5 * a, dtype=tf.float32
                      ),
                      2,
                  ),
              ),
              ("c", lambda b, a: normal.Normal(loc=a, scale=1.0 + b)),
          )
      )
  
      joint = jd_class(models[jd_class], validate_args=True)
  
      z = self.evaluate(joint.sample(seed=test_util.test_seed()))
      a, b, c = z.values() if hasattr(z, "values") else z
  
      log_prob = self.evaluate(joint.log_prob(z))
      prob = self.evaluate(joint.prob(z))
  
      expected_log_prob = self.evaluate(
          np.log(0.5)
          + tf.reduce_sum(
              tf.math.log(
                  b * (0.25 + 0.5 * a) + (1 - b) * (0.75 - 0.5 * a)
              )
          )
          + tf.reduce_sum(
              -0.5 * ((c - a) / (1.0 + b)) ** 2
              - 0.5 * np.log(2.0 * np.pi)
              - tf.math.log((1.0 + b))
          )
      )
      self.assertAllClose(log_prob, expected_log_prob)
      self.assertAllClose(prob, np.exp(expected_log_prob))
  
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 5 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_log_prob_and_prob_with_plate(jd_class) without return types