code: def test_simple_nonlinear_system(self):
      initial_state_prior = mvn_diag.MultivariateNormalDiag(
          0.0, scale_diag=[1.0, 0.3], validate_args=True
      )
      observation_noise_scale = 0.5
  
      # x_{0, t+1} = x_{0, t} - 0.1 * x_{1, t}**3; x_{1, t+1} = x_{1, t}
      def transition_fn(x):
          return mvn_diag.MultivariateNormalDiag(
              tf.stack(
                  [x[..., 0] - 0.1 * tf.pow(x[..., 1], 3), x[..., 1]],
                  axis=-1,
              ),
              scale_diag=[0.5, 0.05],
              validate_args=True,
          )
  
      def transition_jacobian_fn(x):
          return tf.reshape(
              tf.stack(
                  [
                      1.0,
                      -0.3 * x[..., 1] ** 2,
                      tf.zeros(x.shape[:-1]),
                      tf.ones(x.shape[:-1]),
                  ],
                  axis=-1,
              ),
              [2, 2],
          )
  
      def observation_fn(x):
          return mvn_diag.MultivariateNormalDiag(
              x[..., :1],
              scale_diag=[observation_noise_scale],
              validate_args=True,
          )
  
      observation_jacobian_fn = lambda x: [[1.0, 0.0]]
  
      x = [np.zeros((2,), dtype=np.float32)]
      num_timesteps = 20
      for _ in range(num_timesteps - 1):
          x.append(
              transition_fn(x[-1]).sample(seed=test_util.test_seed())
          )
      x = tf.stack(x)
      observations = observation_fn(x).sample(
          seed=test_util.test_seed()
      )
  
      results = extended_kalman_filter.extended_kalman_filter(
          observations=observations,
          initial_state_prior=initial_state_prior,
          transition_fn=transition_fn,
          observation_fn=observation_fn,
          transition_jacobian_fn=transition_jacobian_fn,
          observation_jacobian_fn=observation_jacobian_fn,
      )
      (
          filtered_mean,
          filtered_cov,
          predicted_mean,
          predicted_cov,
          observation_mean,
          observation_cov,
          log_marginal_likelihood,
          timestep,
      ) = results
      self.assertAllEqual(filtered_mean.shape, [num_timesteps, 2])
      self.assertAllEqual(filtered_cov.shape, [num_timesteps, 2, 2])
      self.assertAllEqual(predicted_mean.shape, [num_timesteps, 2])
      self.assertAllEqual(predicted_cov.shape, [num_timesteps, 2, 2])
      self.assertAllEqual(observation_mean.shape, [num_timesteps, 1])
      self.assertAllEqual(observation_cov.shape, [num_timesteps, 1, 1])
      self.assertAllEqual(
          log_marginal_likelihood.shape, [num_timesteps]
      )
      self.assertAllEqual(timestep.shape, [num_timesteps])
  
      # Check that the estimate of the most observable quantity is close to actual
      # (i.e. MSE is close to the observation noise covariance).
      if tf.executing_eagerly():
          # Skip in graph mode, because running the graph resamples the
          # observations.
          self.assertLess(
              tf.reduce_mean((filtered_mean[..., 0] - x[..., 0]) ** 2),
              0.3,
          )
  
      self.assertAllEqual(timestep, np.arange(num_timesteps))
  
      # test structured input
      observations_struct = {
          "a": observations,
          "b": (observations, observations),
      }
      nested_results = self.evaluate(
          extended_kalman_filter.extended_kalman_filter(
              observations=observations_struct,
              initial_state_prior=initial_state_prior,
              transition_fn=transition_fn,
              observation_fn=observation_fn,
              transition_jacobian_fn=transition_jacobian_fn,
              observation_jacobian_fn=observation_jacobian_fn,
          )
      )
  
      if tf.executing_eagerly():
          for result, nested_result in zip(results, nested_results):
              self.assertAllCloseNested(
                  tf.nest.map_structure(
                      lambda _: result, observations_struct
                  ),  # pylint: disable=cell-var-from-loop
                  nested_result,
              )
  
      # test observations with event_size > 1
      def observation_fn_3dim(x):
          loc = tf.stack(
              [
                  x[..., 0],
                  tf.reduce_sum(x, axis=-1),
                  tf.reduce_prod(x, axis=-1),
              ],
              axis=-1,
          )
          return mvn_diag.MultivariateNormalDiag(
              loc, scale_diag=[0.5, 0.5, 0.5], validate_args=True
          )
  
      def observation_jacobian_fn_3dim(x):
          return tf.reshape(
              tf.stack(
                  [1.0, 0.0, 1.0, 1.0, x[..., 1], x[..., 0]], axis=-1
              ),
              [3, 2],
          )
  
      observations_3dim = observation_fn_3dim(x).sample(
          seed=test_util.test_seed()
      )
      results_3dim = extended_kalman_filter.extended_kalman_filter(
          observations=observations_3dim,
          initial_state_prior=initial_state_prior,
          transition_fn=transition_fn,
          observation_fn=observation_fn_3dim,
          transition_jacobian_fn=transition_jacobian_fn,
          observation_jacobian_fn=observation_jacobian_fn_3dim,
      )
  
      self.assertAllEqual(results_3dim[4].shape, [num_timesteps, 3])
      self.assertAllEqual(results_3dim[5].shape, [num_timesteps, 3, 3])
  
file_path: probability/tensorflow_probability/python/experimental/sequential/extended_kalman_filter_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.experimental.sequential import extended_kalman_filter
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the ExtendedKalmanFilterTest obj's underlying member method code with example usage: obj.test_simple_nonlinear_system() and returns: mvn_diag, tf, mvn_diag, mvn_diag, tf