code: def ensemble_kalman_filter_predict(
      state, transition_fn, seed=None, inflate_fn=None, name=None
  ):
      """Ensemble Kalman filter prediction step.
  
      The [Ensemble Kalman Filter](
      https://en.wikipedia.org/wiki/Ensemble_Kalman_filter) is a Monte Carlo
      version of the traditional Kalman Filter. See also [2]. It assumes the model
  
      ```
      X[t] ~ transition_fn(X[t-1])
      Y[t] ~ observation_fn(X[t])
      ```
  
      Given the ensemble `state.particles` sampled from `P(X[t-1] | Y[t-1], ...)`,
      this function produces the predicted (a.k.a. forecast or background) ensemble
      sampled from `P(X[t] | Y[t-1], ...)`. This is the predicted next state
      *before* assimilating the observation `Y[t]`.
  
      Typically, with `F` some deterministic mapping, `transition_fn(X)` returns a
      normal distribution centered at `F(X)`.
  
      Args:
        state: Instance of `EnsembleKalmanFilterState`.
        transition_fn: callable returning a (joint) distribution over the next
          latent state, and any information in the `extra` state.
          Each component should be an instance of
          `MultivariateNormalLinearOperator`.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        inflate_fn: Function that takes in the `particles` and returns a new set of
          `particles`. Used for inflating the covariance of points. Note this
          function should try to preserve the sample mean of the particles, and
          scale up the sample covariance [3].
        name: Python `str` name for ops created by this method.
          Default value: `None` (i.e., `'ensemble_kalman_filter_predict'`).
  
      Returns:
        next_state: `EnsembleKalmanFilterState` representing particles after
          applying `transition_fn`.
  
      #### References
  
      [1] Geir Evensen. Sequential data assimilation with a nonlinear
          quasi-geostrophic model using Monte Carlo methods to forecast error
          statistics. Journal of Geophysical Research, 1994.
  
      [2] Matthias Katzfuss, Jonathan R. Stroud & Christopher K. Wikle
          Understanding the Ensemble Kalman Filter.
          The Americal Statistician, 2016.
  
      [3] Jeffrey L. Anderson and Stephen L. Anderson. A Monte Carlo Implementation
          of the Nonlinear Filtering Problem to Produce Ensemble Assimilations and
          Forecasts. Monthly Weather Review, 1999.
  
      """
      with tf.name_scope(name or "ensemble_kalman_filter_predict"):
          if inflate_fn is not None:
              state = EnsembleKalmanFilterState(
                  step=state.step,
                  particles=inflate_fn(state.particles),
                  extra=state.extra,
              )
  
          new_particles_dist, extra = transition_fn(
              state.step, state.particles, state.extra
          )
  
          return EnsembleKalmanFilterState(
              step=state.step,
              particles=new_particles_dist.sample(seed=seed),
              extra=extra,
          )
  
file_path: probability/tensorflow_probability/python/experimental/sequential/ensemble_kalman_filter.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import mvn_low_rank_update_linear_operator_covariance
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
label_desc: Use image node_id 3 for calling a global function with example usage: ensemble_kalman_filter_predict(state, transition_fn, seed, inflate_fn, name) and returns: EnsembleKalmanFilterState