code: def augment_with_state_history(fn):
      """Decorates a transition or proposal fn to track state history.
  
      For example usage, see
      `tfp.experimental.mcmc.augment_prior_with_state_history`.
  
      Args:
        fn: Python `callable` to wrap, having signature
          `new_state_dist = fn(step, state_with_history, **kwargs)` where
          `state_with_history` is a `StateWithHistory` namedtuple.
      Returns:
        augmented_fn: Python `callable` wrapping `fn`, having signature
          `new_state_with_history_dist = augmented_fn(step, state_with_history,
          **kwargs)`. The return value is a `tfd.JointDistributionNamed` instance
          over`tfp.experimental.mcmc.StateWithHistory` namedtuples, in which the
          `state_history` component is rotated to discard
          the (previously-oldest) state at the initial position and append the
          new state at the final position.
      """
  
      def augmented_fn(step, state_with_history, **kwargs):
          """Builds history-tracking dist. over `StateWithHistory` instances."""
          with tf.name_scope("augment_with_state_history"):
              new_state_dist = fn(step, state_with_history, **kwargs)
  
              def new_state_history_dist(state):
                  with tf.name_scope("new_state_history_dist"):
                      new_state_histories = tf.nest.map_structure(
                          lambda h, s: tf.concat(
                              [
                                  h[
                                      :, 1:
                                  ],  # pylint: disable=g-long-lambda
                                  s[:, tf.newaxis],
                              ],
                              axis=1,
                          ),
                          state_with_history.state_history,
                          state,
                      )
                      return joint_distribution_util.independent_joint_distribution_from_structure(
                          _wrap_as_distributions(new_state_histories)
                      )
  
          return joint_distribution_named.JointDistributionNamed(
              StateWithHistory(
                  state=new_state_dist,
                  state_history=new_state_history_dist,
              )
          )
  
      return augmented_fn
  
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 3 for calling a global function with example usage: augment_with_state_history(fn) and returns: augmented_fn, joint_distribution_named, joint_distribution_util