code: def test_ensemble_kalman_filter_constant_model_multivariate(self):
      def transition_fn(_, particles, extra):
          return (
              mvn_diag.MultivariateNormalDiag(
                  loc=particles, scale_diag=[1e-11] * 2
              ),
              extra,
          )
  
      def observation_fn(_, particles, extra):
          return (
              mvn_diag.MultivariateNormalDiag(
                  loc=particles, scale_diag=[1e-1] * 2
              ),
              extra,
          )
  
      seed_stream = test_util.test_seed_stream()
  
      # Initialize the ensemble.
      particles = self.evaluate(
          tf.random.normal(
              shape=[300, 3, 2], seed=seed_stream(), dtype=tf.float64
          )
      )
  
      state = ekf.EnsembleKalmanFilterState(
          step=0, particles=particles, extra={"unchanged": 1}
      )
  
      for _ in range(8):
          state = ekf.ensemble_kalman_filter_predict(
              state,
              transition_fn=transition_fn,
              seed=seed_stream(),
              inflate_fn=None,
          )
  
          state = ekf.ensemble_kalman_filter_update(
              state,
              observation=[0.0, 0.0],
              observation_fn=observation_fn,
              seed=seed_stream(),
          )
  
      self.assertAllClose(
          [[0.0, 0.0]] * 3,
          self.evaluate(tf.reduce_mean(state.particles, axis=0)),
          atol=1e-2,
      )
  
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 4 for calling the EnsembleKalmanFilterTest obj's underlying member method code with example usage: obj.test_ensemble_kalman_filter_constant_model_multivariate() and returns: mvn_diag,  extra, mvn_diag,  extra