code: def test_broadcast_ragged_batch_shape(self, base_jd_class, jda_class):
      base_jd_models = {}
  
      # Writing a JDC with ragged batch shape will broadcast the first
      # distribution over the second.
      # (though note, this model breaks `log_prob` with nontrivial sample shape).
      def coroutine():
          x = yield Root(normal.Normal(0.0, scale=1.0))
          yield normal.Normal(
              x[..., tf.newaxis], [1.0, 2.0, 3.0, 4.0, 5.0]
          )
  
      base_jd_models[jdc.JointDistributionCoroutine] = coroutine
      base_jd_models[jds.JointDistributionSequential] = [
          normal.Normal(0.0, scale=1.0),
          lambda x: normal.Normal(
              x[..., tf.newaxis], [1.0, 2.0, 3.0, 4.0, 5.0]
          ),
      ]
      base_jd_models[jdn.JointDistributionNamed] = {
          "x": normal.Normal(0.0, scale=1.0),
          "y": lambda x: normal.Normal(
              x[..., tf.newaxis], [1.0, 2.0, 3.0, 4.0, 5.0]
          ),
      }
  
      # But we can get equivalent behavior in a JDCA by expanding dims so that
      # the batch dimensions line up.
      jd_auto_models = {}
  
      def coroutine_auto():
          x = yield normal.Normal(0.0, scale=[1.0])
          yield normal.Normal(x, [1.0, 2.0, 3.0, 4.0, 5.0])
  
      jd_auto_models[
          jdab.JointDistributionCoroutineAutoBatched
      ] = coroutine_auto
      jd_auto_models[jdab.JointDistributionSequentialAutoBatched] = [
          normal.Normal(0.0, scale=[1.0]),
          lambda x: normal.Normal(x, [1.0, 2.0, 3.0, 4.0, 5.0]),
      ]
      jd_auto_models[
          jdab.JointDistributionNamedAutoBatched
      ] = collections.OrderedDict(
          (
              ("x", normal.Normal(0.0, scale=[1.0])),
              (
                  "y",
                  lambda x: normal.Normal(x, [1.0, 2.0, 3.0, 4.0, 5.0]),
              ),
          )
      )
  
      # Writing a JD with ragged batch shape will broadcast the first
      # distribution over the second.
      # (though note, this model breaks `log_prob` with nontrivial sample shape).
      jd_broadcasting = base_jd_class(base_jd_models[base_jd_class])
  
      # This model's broadcasting behavior is a footgun (it can break inference
      # routines and cause silently incorrect optimization); it should be
      # disallowed by `validate_args`.
      with self.assertRaisesRegexp(
          Exception,
          (
              "Component batch shapes are inconsistent|"
              "Broadcasting probably indicates an error in model specification"
          ),
      ):
          jda_invalid = jda_class(
              jd_auto_models[jda_class],
              batch_ndims=1,
              validate_args=True,
          )
          _ = self.evaluate(
              jda_invalid.log_prob(
                  jda_invalid.sample(seed=test_util.test_seed())
              )
          )
  
      # But, if the user wants to run with no guardrails, one can eke out
      # performance wins when evaluating a shared value over multiple models.
      jda_broadcasting = jda_class(
          jd_auto_models[jda_class], batch_ndims=1
      )
  
      self.assertAllEqualNested(
          jda_broadcasting._model_flatten(jda_broadcasting.event_shape),
          [[], []],
      )
      self.assertAllEqual(jda_broadcasting.batch_shape, [5])
  
      joint_sample = jda_broadcasting.sample(seed=test_util.test_seed())
      x_sample, y_sample = self.evaluate(
          list(joint_sample.values())
          if hasattr(joint_sample, "values")
          else joint_sample
      )
      # The model samples only a single value for x, shared across the batch.
      self.assertAllEqual(x_sample.shape, [1])
      self.assertAllEqual(y_sample.shape, [5])
  
      lp_jd_broadcast = self.evaluate(
          jd_broadcasting.log_prob(
              jd_broadcasting._model_unflatten(
                  [x_sample[..., 0], y_sample]
              )
          )
      )
      lp_jda_broadcast = self.evaluate(
          jda_broadcasting.log_prob(
              jda_broadcasting._model_unflatten([x_sample, y_sample])
          )
      )
      self.assertAllEqual(lp_jda_broadcast.shape, [5])
      self.assertAllEqual(lp_jd_broadcast, lp_jda_broadcast)
  
      # Try drawing multiple samples and computing log-prob.
      joint_sample = self.evaluate(
          jda_broadcasting.sample([2, 3], seed=test_util.test_seed())
      )
      lp_jda_broadcast = self.evaluate(
          jda_broadcasting.log_prob(joint_sample)
      )
      self.assertAllEqual(lp_jda_broadcast.shape, [2, 3, 5])
  
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 4 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_broadcast_ragged_batch_shape(base_jd_class, jda_class) without return types