code: def test_ensemble_kalman_filter_constant_univariate_shapes(self):
      # Simple transition model where the state doesn't change,
      # so we are estimating a constant.
  
      def transition_fn(_, particles, extra):
          return (
              mvn_diag.MultivariateNormalDiag(
                  loc=particles, scale_diag=[1e-11]
              ),
              extra,
          )
  
      def observation_fn(_, particles, extra):
          return (
              mvn_diag.MultivariateNormalDiag(
                  loc=particles, scale_diag=[1e-2]
              ),
              extra,
          )
  
      # Initialize the ensemble.
      particles = self.evaluate(
          tf.random.normal(shape=[100, 1], seed=test_util.test_seed())
      )
  
      state = ekf.EnsembleKalmanFilterState(
          step=0, particles=particles, extra={"unchanged": 1}
      )
  
      predicted_state = ekf.ensemble_kalman_filter_predict(
          state,
          transition_fn=transition_fn,
          inflate_fn=None,
          seed=test_util.test_seed(),
      )
  
      observation = tf.convert_to_tensor([0.0], dtype=particles.dtype)
      log_ml = ekf.ensemble_kalman_filter_log_marginal_likelihood(
          predicted_state,
          observation=observation,
          observation_fn=observation_fn,
          seed=test_util.test_seed(),
      )
      self.assertAllEqual(observation.shape[:-1], log_ml.shape)
      log_ml_krazy_obs = (
          ekf.ensemble_kalman_filter_log_marginal_likelihood(
              predicted_state,
              observation=observation + 10.0,
              observation_fn=observation_fn,
              seed=test_util.test_seed(),
          )
      )
      self.assertAllLess(log_ml_krazy_obs, log_ml)
  
      # Check that extra is correctly propagated.
      self.assertIn("unchanged", predicted_state.extra)
      self.assertEqual(1, predicted_state.extra["unchanged"])
  
      # Check that the state respected the constant dynamics.
      self.assertAllClose(state.particles, predicted_state.particles)
  
      updated_state = ekf.ensemble_kalman_filter_update(
          predicted_state,
          # The observation is the constant 0.
          observation=observation,
          seed=test_util.test_seed(),
          observation_fn=observation_fn,
      )
  
      # Check that extra is correctly propagated.
      self.assertIn("unchanged", updated_state.extra)
      self.assertEqual(1, updated_state.extra["unchanged"])
      self.assertAllEqual(
          state.particles.shape, updated_state.particles.shape
      )
  
imports: import collections
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.experimental.sequential import ensemble_kalman_filter
  from tensorflow_probability.python.internal import test_combinations
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 2 for calling the EnsembleKalmanFilterTest obj's underlying member method code with example usage: obj.test_ensemble_kalman_filter_constant_univariate_shapes() and returns: mvn_diag,  extra, mvn_diag,  extra