code: def test_ensemble_kalman_filter_linear_model(self):
      # Simple transition model where the state doesn't change,
      # so we are estimating a constant.
  
      def transition_fn(_, particles, extra):
          particles = {
              "x": particles["x"] + particles["xdot"],
              "xdot": particles["xdot"],
          }
          extra["transition_count"] += 1
          return (
              jdn.JointDistributionNamed(
                  dict(
                      x=mvn_diag.MultivariateNormalDiag(
                          loc=particles["x"], scale_diag=[1e-11]
                      ),
                      xdot=mvn_diag.MultivariateNormalDiag(
                          particles["xdot"], scale_diag=[1e-11]
                      ),
                  )
              ),
              extra,
          )
  
      def observation_fn(_, particles, extra):
          extra["observation_count"] += 1
          return (
              mvn_diag.MultivariateNormalDiag(
                  loc=particles["x"], scale_diag=[1e-2]
              ),
              extra,
          )
  
      seed_stream = test_util.test_seed_stream()
  
      # Initialize the ensemble.
      particles = {
          "x": self.evaluate(
              tf.random.normal(shape=[300, 5, 1], seed=seed_stream())
          ),
          "xdot": self.evaluate(
              tf.random.normal(shape=[300, 5, 1], seed=seed_stream())
          ),
      }
  
      state = ekf.EnsembleKalmanFilterState(
          step=0,
          particles=particles,
          extra={"observation_count": 0, "transition_count": 0},
      )
  
      for i in range(5):
          state = ekf.ensemble_kalman_filter_predict(
              state,
              transition_fn=transition_fn,
              seed=seed_stream(),
              inflate_fn=None,
          )
          self.assertIn("transition_count", state.extra)
          self.assertEqual(i + 1, state.extra["transition_count"])
  
          state = ekf.ensemble_kalman_filter_update(
              state,
              observation=[1.0 * i],
              observation_fn=observation_fn,
              seed=seed_stream(),
          )
  
          self.assertIn("observation_count", state.extra)
          self.assertEqual(i + 1, state.extra["observation_count"])
  
      self.assertAllClose(
          [4.0] * 5,
          self.evaluate(
              tf.reduce_mean(state.particles["x"], axis=[0, -1])
          ),
          rtol=0.05,
      )
  
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 3 for calling the EnsembleKalmanFilterTest obj's underlying member method code with example usage: obj.test_ensemble_kalman_filter_linear_model() and returns: jdn,  extra, mvn_diag,  extra