code: def test_batch_estimation(self):
      # Batch of models, each a normal random walk with unknown scale param.
      batch_shape = [2, 3]
      parameter_prior = (
          transformed_distribution.TransformedDistribution(
              normal.Normal(loc=tf.ones(batch_shape), scale=1.0),
              softplus.Softplus(),
          )
      )
      parameterized_initial_state_prior_fn = (
          lambda parameters: normal.Normal(loc=0.0, scale=parameters)
      )
      parameteried_transition_fn = (
          lambda _, state, parameters: normal.Normal(
              loc=state, scale=parameters
          )
      )
      parameterized_observation_fn = (
          lambda _, state, **kwargs: normal.Normal(loc=state, scale=1.0)
      )
  
      # Generate a batch of synthetic observations from the model.
      num_timesteps = 100
      true_scales = self.evaluate(
          parameter_prior.sample(seed=test_util.test_seed())
      )
      trajectories = tf.math.cumsum(
          tf.random.normal(
              [num_timesteps] + batch_shape, seed=test_util.test_seed()
          )
          * true_scales,
          axis=0,
      )
      observations = self.evaluate(
          parameterized_observation_fn(0, trajectories).sample(
              seed=test_util.test_seed()
          )
      )
  
      # Estimate the batch of scale parameters.
      iterated_filter = iterated_filter_lib.IteratedFilter(
          parameter_prior=parameter_prior,
          parameterized_initial_state_prior_fn=(
              parameterized_initial_state_prior_fn
          ),
          parameterized_transition_fn=parameteried_transition_fn,
          parameterized_observation_fn=parameterized_observation_fn,
      )
      estimated_scales = self.evaluate(
          iterated_filter.estimate_parameters(
              observations=observations,
              num_iterations=20,
              num_particles=1024,
              initial_perturbation_scale=1.0,
              cooling_schedule=(
                  iterated_filter_lib.geometric_cooling_schedule(
                      0.001, k=20
                  )
              ),
              seed=test_util.test_seed(),
          )
      )
      final_scales = tf.nest.map_structure(
          lambda x: x[-1], estimated_scales
      )
      # Note that this inference isn't super precise with the current tuning.
      # Varying the seed, the max absolute error across the batch is typically
      # in the range 0.2 - 0.4.
      self.assertAllClose(
          np.mean(final_scales, axis=0), true_scales, atol=1.0
      )
  
file_path: probability/tensorflow_probability/python/experimental/sequential/iterated_filter_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.sequential import iterated_filter
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the IteratedFilterTest obj's underlying member method code with example usage: obj.test_batch_estimation() without return types