code: def extended_kalman_filter(
      observations,
      initial_state_prior,
      transition_fn,
      observation_fn,
      transition_jacobian_fn,
      observation_jacobian_fn,
      name=None,
  ):
      """Applies an Extended Kalman Filter to observed data.
  
      The [Extended Kalman Filter](
      https://en.wikipedia.org/wiki/Extended_Kalman_filter) is a nonlinear version
      of the Kalman filter, in which the transition function is linearized by
      first-order Taylor expansion around the current mean and covariance of the
      state estimate.
  
      Args:
        observations: a (structure of) `Tensor`s, each of shape
          `concat([[num_timesteps, b1, ..., bN], [event_size]])` with scalar
          `event_size` and optional batch dimensions `b1, ..., bN`.
        initial_state_prior: a `tfd.Distribution` instance (typically
          `MultivariateNormal`) with `event_shape` equal to `state_size` and an
          optional `batch_shape` of [`b1, ..., bN`], representing the prior over the
          state.
        transition_fn: a Python `callable` that accepts (batched) vectors of length
          `state_size`, and returns a `tfd.Distribution` instance, typically a
          `MultivariateNormal`, representing the state transition and covariance.
        observation_fn: a Python `callable` that accepts a (batched) vector of
          length `state_size` and returns a `tfd.Distribution` instance, typically
          a `MultivariateNormal` representing the observation model and covariance.
        transition_jacobian_fn: a Python `callable` that accepts a (batched) vector
          of length `state_size` and returns a (batched) matrix of shape
          `[state_size, state_size]`, representing the Jacobian of `transition_fn`.
        observation_jacobian_fn: a Python `callable` that accepts a (batched) vector
          of length `state_size` and returns a (batched) matrix of size
          `[state_size, event_size]`, representing the Jacobian of `observation_fn`.
        name: Python `str` name for ops created by this method.
          Default value: `None` (i.e., `'extended_kalman_filter'`).
      Returns:
        filtered_mean: a (structure of) `Tensor`(s) of shape
          `concat([[num_timesteps, b1, ..., bN], [state_size]])`. The mean of the
          filtered state estimate.
        filtered_cov: a (structure of) `Tensor`(s) of shape
          `concat([[num_timesteps, b1, ..., bN], [state_size, state_size]])`.
           The covariance of the filtered state estimate.
        predicted_mean: a (structure of) `Tensor`(s) of shape
          `concat([[num_timesteps, b1, ..., bN], [state_size]])`. The prior
          predicted means of the state.
        predicted_cov: a (structure of) `Tensor`(s) of shape
          `concat([[num_timesteps, b1, ..., bN], [state_size, state_size]])`
          The prior predicted covariances of the state estimate.
        observation_mean: a (structure of) `Tensor`(s) of shape
          `concat([[num_timesteps, b1, ..., bN], [event_size]])`. The prior
          predicted mean of observations.
        observation_cov: a (structure of) `Tensor`(s) of shape
          `concat([[num_timesteps, b1, ..., bN], [event_size, event_size]])`. The
          prior predicted covariance of observations.
        log_marginal_likelihood: a (structure of) `Tensor`(s) of shape
          `[num_timesteps, b1, ..., bN]`. Log likelihood of the observations with
          respect to the observation.
        timestep: a (structure of) integer `Tensor`(s) of shape
          `[num_timesteps, b1, ..., bN]` containing time indices.
  
      #### Examples
  
      **Estimate a simple nonlinear system**: Let's consider a system defined by
      the transition equation `y_{t+1} = y_t - 0.1 * w_t **3` and `w_{t+1} = w_t`,
      such that the state can be expressed as `[y, w]`. The `transition_fn` and
      `transition_jacobian_fn` can be expressed as:
  
      ```python
      def transition_fn(x):
        return tfd.MultivariateNormalDiag(
            tf.stack(
                [x[..., 0] - 0.1 * x[..., 1]**3, x[..., 1]], axis=-1),
            scale_diag=[0.7, 0.2])
  
      def transition_jacobian_fn(x):
        return tf.reshape(
          tf.stack(
              [1. - 0.1 * x[..., 1]**3, -0.3 * x[..., 1]**2,
              tf.zeros(x.shape[:-1]), tf.ones(x.shape[:-1])], axis=-1),
          [2, 2])
      ```
  
      Assume we take noisy measurements of only the first element of the state.
  
      ```python
      observation_fn = lambda x: tfd.MultivariateNormalDiag(
          x[..., :1], scale_diag=[1.])
      observation_jacobian_fn = lambda x: [[1., 0.]]
      ```
  
      We define a prior over the initial state, and use it to synthesize data for
      20 steps of the process.
  
      ```python
      initial_state_prior = tfd.MultivariateNormalDiag(0., scale_diag=[1., 0.3])
  
      x = [np.zeros((2,), dtype=np.float32)]
      for t in range(20):
        x.append(transition_fn(x[-1]).sample())
      x = tf.stack(x)
  
      observations=observation_fn(x).sample()
      ```
  
      Run the Extended Kalman filter on the synthesized observed data.
  
      ```python
      results = tfp.experimental.sequential.extended_kalman_filter(
          observations,
          initial_state_prior,
          transition_fn,
          observation_fn,
          transition_jacobian_fn,
          observation_jacobian_fn)
      ```
      """
      with tf.name_scope(name or "extended_kalman_filter"):
          observations_shape = prefer_static.shape(
              tf.nest.flatten(observations)[0]
          )
          num_timesteps = observations_shape[0]
  
          # add singleton batch dimensions to initial state estimates, if necessary
          dummy_zeros = tf.zeros(observations_shape[1:-1])
          initial_state = (
              initial_state_prior.mean() + dummy_zeros[..., tf.newaxis]
          )
          initial_covariance = (
              initial_state_prior.covariance()
              + dummy_zeros[..., tf.newaxis, tf.newaxis]
          )
          observation_dist = observation_fn(initial_state)
  
          # Initialize the state estimate.
          def _make_initial_kalman_filter_state():
              return linear_gaussian_ssm.KalmanFilterState(
                  predicted_mean=initial_state,
                  predicted_cov=initial_covariance,
                  filtered_mean=initial_state,
                  filtered_cov=initial_covariance,
                  observation_mean=observation_dist.mean(),
                  observation_cov=observation_dist.covariance(),
                  log_marginal_likelihood=dummy_zeros,
                  timestep=tf.zeros(
                      observations_shape[1:-1], dtype=tf.int32
                  )
                  - 1,
              )
  
          initial_estimate = tf.nest.map_structure(
              lambda _: _make_initial_kalman_filter_state(),
              observations,
          )
  
          initial_accumulated_quantities = (
              _initialize_accumulated_quantities(
                  _make_initial_kalman_filter_state(),
                  observations,
                  num_timesteps,
              )
          )
  
          run_ekf_step = functools.partial(
              extended_kalman_filter_one_step,
              transition_fn=transition_fn,
              observation_fn=observation_fn,
              transition_jacobian_fn=transition_jacobian_fn,
              observation_jacobian_fn=observation_jacobian_fn,
          )
  
          def _loop_body(
              step, current_estimate, accumulated_quantities
          ):
              """Body of the `while_loop` for running forward filtering."""
              current_observations = tf.nest.map_structure(
                  lambda x, step=step: tf.gather(x, step), observations
              )
              updated_estimate = nest.map_structure_up_to(
                  current_observations,
                  run_ekf_step,
                  current_estimate,
                  current_observations,
              )
              new_accumulated_quantities = (
                  _write_accumulated_quantities(
                      step, accumulated_quantities, updated_estimate
                  )
              )
              return (
                  step + 1,
                  updated_estimate,
                  new_accumulated_quantities,
              )
  
          _, _, loop_results = tf.while_loop(
              cond=lambda step, *_: step < num_timesteps,
              body=_loop_body,
              loop_vars=(
                  tf.zeros([], tf.int32),
                  initial_estimate,
                  initial_accumulated_quantities,
              ),
          )
          return [
              tf.nest.map_structure(lambda ta: ta.stack(), x)
              for x in loop_results
          ]
  
file_path: probability/tensorflow_probability/python/experimental/sequential/extended_kalman_filter.py
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math import linalg
  from tensorflow.python.util import nest
label_desc: Use image node_id 4 for calling a global function with example usage: extended_kalman_filter(observations, initial_state_prior, transition_fn, observation_fn, transition_jacobian_fn, observation_jacobian_fn, name) and returns: unknown, linear_gaussian_ssm, unknown,  updated_estimate,  new_accumulated_quantities