code: def test_model_can_use_state_history(self):
      initial_state_prior = (
          particle_filter_augmentation.augment_prior_with_state_history(
              jdn.JointDistributionNamed({"x": poisson.Poisson(1.0)}),
              history_size=2,
          )
      )
  
      # Deterministic dynamics compute a Fibonacci sequence.
      @particle_filter_augmentation.augment_with_state_history
      def fibonacci_transition_fn(step, state_with_history):
          del step
          return jdn.JointDistributionNamed(
              {
                  "x": deterministic.Deterministic(
                      tf.reduce_sum(
                          state_with_history.state_history["x"][
                              ..., -2:
                          ],
                          axis=-1,
                      )
                  )
              }
          )
  
      # We'll observe the ratio of the current and previous state.
      def observe_ratio_of_last_two_states_fn(_, state_with_history):
          ratio = tf.ones_like(state_with_history.state["x"])
          if state_with_history.state_history is not None:
              ratio = state_with_history.state["x"] / (
                  state_with_history.state_history["x"][..., -2] + 1e-6
              )  # Avoid division by 0.
          return normal.Normal(loc=ratio, scale=0.1)
  
      # The ratios between successive terms of a Fibonacci sequence
      # should, in the limit, approach the golden ratio.
      golden_ratio = (1.0 + np.sqrt(5.0)) / 2.0
      observed_ratios = np.array([golden_ratio] * 10).astype(self.dtype)
  
      trajectories_with_history, lps = self.evaluate(
          particle_filter.infer_trajectories(
              observed_ratios,
              initial_state_prior=initial_state_prior,
              transition_fn=fibonacci_transition_fn,
              observation_fn=observe_ratio_of_last_two_states_fn,
              num_particles=100,
              seed=test_util.test_seed(),
          )
      )
      trajectories = trajectories_with_history.state
  
      # Verify that we actually produced Fibonnaci sequences.
      self.assertAllClose(
          trajectories["x"][2:],
          trajectories["x"][1:-1] + trajectories["x"][:-2],
      )
  
      # Ratios should get closer to golden as the series progresses, so
      # likelihoods will increase.
      self.assertAllGreaterEqual(lps[2:] - lps[:-2], 0.0)
  
      # Any particles that sampled initial values of 0. should have been
      # discarded, since those lead to degenerate series that do not approach
      # the golden ratio.
      self.assertAllGreaterEqual(trajectories["x"][0], 1.0)
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/particle_filter_augmentation_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.experimental.mcmc import particle_filter
  from tensorflow_probability.python.experimental.mcmc import particle_filter_augmentation
  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 _ParticleFilterUtilTest obj's underlying member method code with example usage: obj.test_model_can_use_state_history() and returns: jdn, normal