code: def test_warns_if_transition_distribution_has_unexpected_shape(self):
      initial_state_prior = jdab.JointDistributionNamedAutoBatched(
          {
              "sales": deterministic.Deterministic(0.0),
              "inventory": deterministic.Deterministic(1000.0),
          }
      )
  
      # Inventory decreases by a Poisson RV 'sales', but is lower bounded at zero.
      def valid_transition_fn(_, particles):
          return jdab.JointDistributionNamedAutoBatched(
              {
                  "sales": poisson.Poisson(
                      10.0 * tf.ones_like(particles["inventory"])
                  ),
                  "inventory": lambda sales: deterministic.Deterministic(  # pylint: disable=g-long-lambda
                      tf.maximum(0.0, particles["inventory"] - sales)
                  ),
              },
              batch_ndims=1,
              validate_args=True,
          )
  
      def dummy_observation_fn(_, state):
          return normal.Normal(state["inventory"], 1000.0)
  
      run_filter = functools.partial(
          particle_filter.particle_filter,
          observations=tf.zeros([10]),
          initial_state_prior=initial_state_prior,
          observation_fn=dummy_observation_fn,
          num_particles=3,
          seed=test_util.test_seed(sampler_type="stateless"),
      )
  
      # Check that the model runs as written.
      self.evaluate(run_filter(transition_fn=valid_transition_fn))
      self.evaluate(
          run_filter(
              transition_fn=valid_transition_fn,
              proposal_fn=valid_transition_fn,
          )
      )
  
      # Check that broken transition functions raise exceptions.
      def transition_fn_broadcasts_over_particles(_, particles):
          return jdn.JointDistributionNamed(
              {
                  "sales": poisson.Poisson(
                      10.0
                  ),  # Proposes same value for all particles.
                  "inventory": lambda sales: deterministic.Deterministic(  # pylint: disable=g-long-lambda
                      tf.maximum(0.0, particles["inventory"] - sales)
                  ),
              },
              validate_args=True,
          )
  
      def transition_fn_partial_batch_shape(_, particles):
          return jdn.JointDistributionNamed(
              # Using `Sample` ensures iid proposals for each particle, but not
              # per-particle log probs.
              {
                  "sales": sample_dist_lib.Sample(
                      poisson.Poisson(10.0),
                      ps.shape(particles["sales"]),
                  ),
                  "inventory": lambda sales: deterministic.Deterministic(  # pylint: disable=g-long-lambda
                      tf.maximum(0.0, particles["inventory"] - sales)
                  ),
              },
              validate_args=True,
          )
  
      def transition_fn_no_batch_shape(_, particles):
          # Autobatched JD defaults to treating num_particles as event shape, but
          # we need it to be batch shape to get per-particle logprobs.
          return jdab.JointDistributionNamedAutoBatched(
              {
                  "sales": poisson.Poisson(
                      10.0 * tf.ones_like(particles["inventory"])
                  ),
                  "inventory": lambda sales: deterministic.Deterministic(  # pylint: disable=g-long-lambda
                      tf.maximum(0.0, particles["inventory"] - sales)
                  ),
              },
              validate_args=True,
          )
  
      with self.assertRaisesRegex(
          ValueError, "transition distribution"
      ):
          self.evaluate(
              run_filter(
                  transition_fn=transition_fn_broadcasts_over_particles
              )
          )
      with self.assertRaisesRegex(
          ValueError, "transition distribution"
      ):
          self.evaluate(
              run_filter(
                  transition_fn=transition_fn_partial_batch_shape
              )
          )
      with self.assertRaisesRegex(
          ValueError, "transition distribution"
      ):
          self.evaluate(
              run_filter(transition_fn=transition_fn_no_batch_shape)
          )
  
      with self.assertRaisesRegex(ValueError, "proposal distribution"):
          self.evaluate(
              run_filter(
                  transition_fn=valid_transition_fn,
                  proposal_fn=transition_fn_partial_batch_shape,
              )
          )
      with self.assertRaisesRegex(ValueError, "proposal distribution"):
          self.evaluate(
              run_filter(
                  transition_fn=valid_transition_fn,
                  proposal_fn=transition_fn_broadcasts_over_particles,
              )
          )
  
      with self.assertRaisesRegex(ValueError, "proposal distribution"):
          self.evaluate(
              run_filter(
                  transition_fn=valid_transition_fn,
                  proposal_fn=transition_fn_no_batch_shape,
              )
          )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  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 transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.mcmc import particle_filter
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 12 for calling the _ParticleFilterTest obj's underlying member method code with example usage: obj.test_warns_if_transition_distribution_has_unexpected_shape() and returns: jdab, normal, jdn, jdn, jdab