code: def extended_kalman_filter_one_step(
      state,
      observation,
      transition_fn,
      observation_fn,
      transition_jacobian_fn,
      observation_jacobian_fn,
      name=None,
  ):
      """A single step of the EKF.
  
      Args:
        state: A `Tensor` of shape
          `concat([[num_timesteps, b1, ..., bN], [state_size]])` with scalar
          `event_size` and optional batch dimensions `b1, ..., bN`.
        observation: A `Tensor` of shape
          `concat([[num_timesteps, b1, ..., bN], [event_size]])` with scalar
          `event_size` and optional batch dimensions `b1, ..., bN`.
        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_one_step'`).
      Returns:
        updated_state: `KalmanFilterState` object containing the updated state
          estimate.
      """
      with tf.name_scope(
          name or "extended_kalman_filter_one_step"
      ) as name:
          # If observations are scalar, we can avoid some matrix ops.
          observation_size_is_static_and_scalar = (
              observation.shape[-1] == 1
          )
  
          current_state = state.filtered_mean
          current_covariance = state.filtered_cov
          current_jacobian = transition_jacobian_fn(current_state)
          state_prior = transition_fn(current_state)
  
          predicted_cov = (
              tf.matmul(
                  current_jacobian,
                  tf.matmul(
                      current_covariance,
                      current_jacobian,
                      transpose_b=True,
                  ),
              )
              + state_prior.covariance()
          )
          predicted_mean = state_prior.mean()
  
          observation_dist = observation_fn(predicted_mean)
          observation_mean = observation_dist.mean()
          observation_cov = observation_dist.covariance()
  
          predicted_jacobian = observation_jacobian_fn(predicted_mean)
          tmp_obs_cov = tf.matmul(predicted_jacobian, predicted_cov)
          residual_covariance = (
              tf.matmul(
                  predicted_jacobian, tmp_obs_cov, transpose_b=True
              )
              + observation_cov
          )
  
          if observation_size_is_static_and_scalar:
              gain_transpose = tmp_obs_cov / residual_covariance
          else:
              chol_residual_cov = tf.linalg.cholesky(
                  residual_covariance
              )
              gain_transpose = linalg.hpsd_solve(
                  residual_covariance,
                  tmp_obs_cov,
                  cholesky_matrix=chol_residual_cov,
              )
  
          filtered_mean = (
              predicted_mean
              + tf.matmul(
                  gain_transpose,
                  (observation - observation_mean)[..., tf.newaxis],
                  transpose_a=True,
              )[..., 0]
          )
  
          tmp_term = -tf.matmul(
              predicted_jacobian, gain_transpose, transpose_a=True
          )
          tmp_term = tf.linalg.set_diag(
              tmp_term, tf.linalg.diag_part(tmp_term) + 1.0
          )
          filtered_cov = tf.matmul(
              tmp_term,
              tf.matmul(predicted_cov, tmp_term),
              transpose_a=True,
          ) + tf.matmul(
              gain_transpose,
              tf.matmul(observation_cov, gain_transpose),
              transpose_a=True,
          )
  
          if observation_size_is_static_and_scalar:
              # A plain Normal would have event shape `[]`; wrapping with Independent
              # ensures `event_shape=[1]` as required.
              predictive_dist = independent.Independent(
                  normal.Normal(
                      loc=observation_mean,
                      scale=tf.sqrt(residual_covariance[..., 0]),
                  ),
                  reinterpreted_batch_ndims=1,
              )
  
          else:
              predictive_dist = mvn_tril.MultivariateNormalTriL(
                  loc=observation_mean, scale_tril=chol_residual_cov
              )
  
          log_marginal_likelihood = predictive_dist.log_prob(
              observation
          )
  
          return linear_gaussian_ssm.KalmanFilterState(
              filtered_mean=filtered_mean,
              filtered_cov=filtered_cov,
              predicted_mean=predicted_mean,
              predicted_cov=predicted_cov,
              observation_mean=observation_mean,
              observation_cov=observation_cov,
              log_marginal_likelihood=log_marginal_likelihood,
              timestep=state.timestep + 1,
          )
  
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 5 for calling a global function with example usage: extended_kalman_filter_one_step(state, observation, transition_fn, observation_fn, transition_jacobian_fn, observation_jacobian_fn, name) and returns: linear_gaussian_ssm