code: def augment_with_observation_history(
      observations, history_size, num_transitions_per_observation=1
  ):
      """Decorates a function to take `observation_history`.
  
      Args:
        observations: a (structure of) Tensors, each of shape
          `concat([[num_observation_steps, b1, ..., bN], event_shape])` with
          optional batch dimensions `b1, ..., bN`.
        history_size: integer `Tensor` number of steps of history to pass.
        num_transitions_per_observation: integer `Tensor` number of
          state transitions between regular observation points. A value of `1`
          indicates that there is an observation at every timestep,
          `2` that every other step is observed, and so on. Values greater than `1`
          may be used with an appropriately-chosen transition function to
          approximate continuous-time dynamics. The initial and final steps
          (steps `0` and `num_timesteps - 1`) are always observed.
          Default value: `1`.
      Returns:
        augment_fn: Python `callable` such that `augmented_fn = augment_fn(fn)`.
          When called, `augmented_fn` invokes `fn`
          with an additional `observation_history` keyword arg, whose value is a
          `Tensor` of shape `concat([[history_size, b1, ..., bN], event_shape])`
          containing up to the most recent `history_size` observations.
      """
  
      def augment_fn(fn):
          """Set the `observation_history` kwarg of the given fn."""
  
          def augmented_fn(step, *args, **kwargs):
              with tf.name_scope("augment_with_observation_history"):
                  observation_idx = (
                      step // num_transitions_per_observation
                  )
                  observation_history_indices = ps.range(
                      ps.maximum(0, observation_idx - history_size),
                      observation_idx,
                  )
                  return fn(
                      step,
                      *args,
                      observation_history=tf.gather(
                          observations, observation_history_indices
                      ),
                      **kwargs,
                  )
  
          return augmented_fn
  
      return augment_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 1 for calling a global function with example usage: augment_with_observation_history(observations, history_size, num_transitions_per_observation) and returns: augment_fn, augmented_fn, fn