code: def _gaussian_momentum_refresh_fn(
      old_momentum: fun_mc.State,
      damping: Optional[float | fun_mc.FloatArray] = 0.0,
      step_size: Optional[float | fun_mc.FloatArray] = 1.0,
      named_axis: Optional[fun_mc.StringNest] = None,
      seed: Optional[Any] = None,
  ) -> fun_mc.State:
      """Momentum refresh function for Gaussian momentum distribution."""
      if named_axis is None:
          named_axis = util.map_tree(lambda _: [], old_momentum)
      damping = fun_mc.maybe_broadcast_structure(damping, old_momentum)
      step_size = fun_mc.maybe_broadcast_structure(
          step_size, old_momentum
      )
      decay_fraction = util.map_tree(
          lambda d, s: jnp.exp(-d * s), damping, step_size
      )
      noise_fraction = util.map_tree(
          lambda df: jnp.sqrt(1.0 - jnp.square(df)), decay_fraction
      )
  
      def _sample_part(
          old_momentum, seed, named_axis, decay_fraction, noise_fraction
      ):
          seed = backend.distribute_lib.fold_in_axis_index(
              seed, named_axis
          )
          return (
              decay_fraction * old_momentum
              + noise_fraction
              * util.random_normal(
                  old_momentum.shape, old_momentum.dtype, seed
              )
          )
  
      seeds = util.unflatten_tree(
          old_momentum,
          util.split_seed(seed, len(util.flatten_tree(old_momentum))),
      )
      new_momentum = util.map_tree_up_to(
          old_momentum,
          _sample_part,
          old_momentum,
          seeds,
          named_axis,
          decay_fraction,
          noise_fraction,
      )
      return new_momentum
  
file_path: probability/spinoffs/fun_mc/fun_mc/malt.py
imports: import functools
  from typing import Any, Callable, NamedTuple, Optional, Tuple
  from fun_mc import backend
  from fun_mc import fun_mc_lib
label_desc: Use image node_id 1 for calling a global function with example usage: _gaussian_momentum_refresh_fn(old_momentum, damping, step_size, named_axis, seed) and returns: new_momentum, unknown