code: def augment_prior_with_state_history(prior, history_size):
      """Augments a prior or proposal distribution's state space with history.
  
      The augmented state space is over `tfp.experimental.mcmc.StateWithHistory`
      namedtuples, which contain the original `state` as well as a `state_history`.
      The `state_history` is a structure of `Tensor`s matching `state`, of shape
      `concat([[num_particles, history_size], state.shape[1:]])`. In other words,
      previous states for each particle are indexed along `axis=1`, to the right
      of the particle indices.
  
      Args:
        prior: a (joint) distribution over the initial latent state,
          with optional batch shape `[b1, ..., bN]`.
        history_size: integer `Tensor` number of steps of history to pass.
      Returns:
        augmented_prior: a `tfd.JointDistributionNamed` instance whose samples
          are `tfp.experimental.mcmc.StateWithHistory` namedtuples.
  
      #### Example
  
      As a toy example, let's see how we'd use state history to experiment with
      stochastic 'Fibonacci sequences'. We'll assume that the sequence starts at a
      value sampled from a Poisson distribution.
  
      ```python
      initial_state_prior = tfd.Poisson(5.)
      initial_state_with_history_prior = (
        tfp.experimental.mcmc.augment_prior_with_state_history(
          initial_state_prior, history_size=2))
      ```
  
      Note that we've augmented the state space to include a state history of
      size two. The augmented state space is over instances of
      `tfp.experimental.mcmc.StateWithHistory`. Initially, the state history
      will simply tile the initial state: if
      `s = initial_state_with_history_prior.sample()`, then
      `s.state_history==[s.state, s.state]`.
  
      Next, we'll define a `transition_fn` that uses the history to
      sample the next integer in the sequence, also from a Poisson distribution.
  
      ```python
      @tfp.experimental.mcmc.augment_with_state_history
      def fibonacci_transition_fn(_, state_with_history):
        expected_next_element = tf.reduce_sum(
          state_with_history.state_history[:, -2:], axis=1)
        return tfd.Poisson(rate=expected_next_element)
      ```
  
      Our transition function must accept `state_with_history`,
      so that it can access the history, but it returns a distribution
      only over the next state. Decorating it with `augment_with_state_history`
      ensures that the state history is automatically propagated.
  
      Note: if we were using an `initial_state_proposal` and/or `proposal_fn`, we
      would need to wrap them similarly to the prior and transition function
      shown here.
  
      Combined with an observation function (which must also now be defined on the
      augmented `StateWithHistory` space), we can track stochastic Fibonacci
      sequences and, for example, infer the initial value of a sequence:
  
      ```python
  
      def observation_fn(_, state_with_history):
        return tfd.Poisson(rate=state_with_history.state)
  
      trajectories, _ = tfp.experimental.mcmc.infer_trajectories(
        observations=tf.convert_to_tensor([4., 11., 16., 23., 40., 69., 100.]),
        initial_state_prior=initial_state_with_history_prior,
        transition_fn=fibonacci_transition_fn,
        observation_fn=observation_fn,
        num_particles=1024)
      inferred_initial_states = trajectories.state[0]
      print(tf.unique_with_counts(inferred_initial_states))
      ```
  
      """
  
      def initialize_state_history(state):
          """Build an initial state history by replicating the initial state."""
          with tf.name_scope("initialize_state_history"):
              initial_state_histories = tf.nest.map_structure(
                  lambda x: tf.broadcast_to(  # pylint: disable=g-long-lambda
                      tf.expand_dims(x, ps.minimum(ps.rank(x), 1)),
                      ps.concat(
                          [
                              ps.shape(x)[:1],
                              [history_size],
                              ps.shape(x)[1:],
                          ],
                          axis=0,
                      ),
                  ),
                  state,
              )
              return joint_distribution_util.independent_joint_distribution_from_structure(
                  _wrap_as_distributions(initial_state_histories)
              )
  
      return joint_distribution_named.JointDistributionNamed(
          StateWithHistory(
              state=prior, state_history=initialize_state_history
          )
      )
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/particle_filter_augmentation.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 4 for calling a global function with example usage: augment_prior_with_state_history(prior, history_size) and returns: joint_distribution_named, joint_distribution_util