code: def persistent_hamiltonian_monte_carlo_step(
      phmc_state: PersistentHamiltonianMonteCarloState,
      target_log_prob_fn: fun_mc.PotentialFn,
      step_size: Optional[Any] = None,
      num_integrator_steps: Optional[fun_mc.IntArray] = None,
      noise_fraction: Optional[fun_mc.FloatArray] = None,
      mh_drift: Optional[fun_mc.FloatArray] = None,
      kinetic_energy_fn: Optional[fun_mc.PotentialFn] = None,
      momentum_sample_fn: Optional[PersistentMomentumSampleFn] = None,
      integrator_trace_fn: Callable[
          [fun_mc.IntegratorStepState, fun_mc.IntegratorStepExtras],
          fun_mc.ArrayNest,
      ] = lambda *args: (),
      log_uniform: Optional[fun_mc.FloatArray] = None,
      integrator_fn: Optional[
          Callable[
              [fun_mc.IntegratorState, fun_mc.FloatArray],
              tuple[fun_mc.IntegratorState, fun_mc.IntegratorExtras],
          ]
      ] = None,
      unroll_integrator: bool = False,
      max_num_integrator_steps: Optional[fun_mc.IntArray] = None,
      energy_change_fn: Callable[
          [
              fun_mc.IntegratorState,
              fun_mc.IntegratorState,
              fun_mc.IntegratorExtras,
          ],
          tuple[fun_mc.FloatArray, Any],
      ] = (
          fun_mc._default_hamiltonian_monte_carlo_energy_change_fn  # pylint: disable=protected-access
      ),
      named_axis: Optional[fun_mc.StringNest] = None,
      seed=None,
  ) -> tuple[
      PersistentHamiltonianMonteCarloState,
      PersistentHamiltonianMonteCarloExtra,
  ]:
      """A step of the Persistent Hamiltonian Monte Carlo `TransitionOperator`.
  
      This is an implementation of the generalized HMC with persistent momentum
      described in [1] (algorithm 15) combined with the persistent Metropolis
      Hastings test from [2]. This generalizes the regular HMC with persistent
      momentum from [3] and the various underdamped langevin dynamics schemes (e.g.
      [4]).
  
      The generalization lies in the free choice of `momentum_sample_fn` and
      `kinetic_energy_fn`. The former forms a Markov Chain with the stationary
      distribution implied by the `kinetic_energy_fn`. By default, the standard
      quadratic kinetic energy is used and the underdamped update is used for
      `momentum_sample_fn`, namely:
  
      ```none
      new_momentum = (
          (1 - noise_fraction**2)**0.5 * old_momentum  +
          noise_fraction * eps)
      eps ~ Normal(0, 1)
      ```
  
      Here are the parameter settings for few special cases:
  
      1. Persistent Hamiltonian Monte Carlo [1] + persistent MH [2]:
  
      ```none
      num_integrator_steps >= 1
      step_size > 0
      noise_fraction in [0, 1]
      mh_drift = 0.03
      ```
  
      Empirical results suggest that if `num_integrator_steps == 1`, then a
      reasonable value for `mh_drift` is `1 - (1 - noise_fraction**2)**0.5`.
  
      2. Unadjusted Underdamped Langevin Dynamics (see [4]):
  
      ```none
      num_integrator_steps = 1
      step_size > 0
      noise_fraction = (1 - exp(-2 * step_size * damping))**0.5
      # This disables the MH step for all but most extreme divergences.
      log_uniform = -1000
      ```
  
      `damping` refers to the parameter in the SDE formulation of the algorithm:
  
      ```none
      dv_t = -damping * v_t * dt - grad(f)(x_t) * dt + (2 * damping)**0.5 * dB_t
      dx_t = v_t * dt
      ```
  
      Args:
        phmc_state: `PersistentHamiltonianMonteCarloState`.
        target_log_prob_fn: Target log prob fn.
        step_size: Step size, structure broadcastable to the `target_log_prob_fn`
          state. Optional if `integrator_fn` is specified.
        num_integrator_steps: Number of integrator steps to take. Optional if
          `integrator_fn` is specified.
        noise_fraction: Noise fraction when refreshing momentum. Optional if
          `momentum_sample_fn` is specified.
        mh_drift: Metropolis Hastings drift term. Optional if `log_uniform` is
          specified.
        kinetic_energy_fn: Kinetic energy function.
        momentum_sample_fn: Sampler for the momentum.
        integrator_trace_fn: Trace function for the integrator.
        log_uniform: Optional logarithm of a uniformly distributed random sample in
          [0, 1], used for the MH accept/reject step.
        integrator_fn: Integrator to use for the HMC dynamics. Uses a
          `hamiltonian_integrator` with `leapfrog_step` by default.
        unroll_integrator: Whether to unroll the loop in the integrator. Only works
          if `num_integrator_steps`/`max_num_integrator_steps' is statically known.
          Ignored if `integrator_fn` is specified.
        max_num_integrator_steps: Maximum number of integrator steps to take. Useful
          when `num_integrator_steps` is dynamic, and yet you still want
          gradients/tracing to work. Ignored if `integrator_fn` is specified.
        energy_change_fn: Callable with signature: `(current_integrator_state,
          proposed_integrator_state,) -> (energy_change, energy_change_extra)`.
          Computes the change in energy between current and proposed states. By
          default, it just substracts the current and proposed energies. A typical
          reason to override this is to improve numerical stability.
        named_axis: Named axes of the state, same structure as `hmc_state.state`.
        seed: For reproducibility.
  
      Returns:
        phmc_state: PersistentHamiltonianMonteCarloState
        phmc_extra: PersistentHamiltonianMonteCarloExtra
  
      #### References
  
      [1]: Neklyudov, K., Welling, M., Egorov, E., & Vetrov, D. (2020). Involutive
           MCMC: a Unifying Framework.
  
      [2]: Neal, R. M. (2020). Non-reversibly updating a uniform [0,1] value for
           Metropolis accept/reject decisions.
  
      [3]: Horowitz, A. M. (1991). A generalized guided Monte Carlo algorithm.
           Physics Letters. [Part B], 268(2), 247-252.
  
      [4]: Ma, Y.-A., Chatterji, N., Cheng, X., Flammarion, N., Bartlett, P., &
           Jordan, M. I. (2019). Is There an Analog of Nesterov Acceleration for
           MCMC?
      """
      state = phmc_state.state
      momentum = phmc_state.momentum
      direction = phmc_state.direction
      state_grads = phmc_state.state_grads
      target_log_prob = phmc_state.target_log_prob
      state_extra = phmc_state.state_extra
      pmh_state = phmc_state.pmh_state
  
      # Impute the optional args.
      if kinetic_energy_fn is None:
          kinetic_energy_fn = fun_mc.make_gaussian_kinetic_energy_fn(
              len(target_log_prob.shape), named_axis=named_axis
          )
  
      if momentum_sample_fn is None:
          if named_axis is None:
              named_axis = util.map_tree(lambda _: [], state)
  
          def _momentum_sample_fn(
              old_momentum: fun_mc.State, seed: Any
          ) -> tuple[fun_mc.State, tuple[()]]:
              seeds = util.unflatten_tree(
                  old_momentum,
                  util.split_seed(
                      seed, len(util.flatten_tree(old_momentum))
                  ),
              )
  
              def _sample_part(old_momentum, seed, named_axis):
                  seed = backend.distribute_lib.fold_in_axis_index(
                      seed, named_axis
                  )
                  return jnp.sqrt(
                      1 - jnp.square(noise_fraction)
                  ) * old_momentum + noise_fraction * util.random_normal(
                      old_momentum.shape, old_momentum.dtype, seed
                  )
  
              new_momentum = util.map_tree_up_to(
                  state, _sample_part, old_momentum, seeds, named_axis
              )
              return new_momentum
  
          momentum_sample_fn = _momentum_sample_fn
  
      if integrator_fn is None:
          step_size = util.map_tree(jnp.asarray, step_size)
          step_size = fun_mc.maybe_broadcast_structure(step_size, state)
  
          def _integrator_fn(
              state: fun_mc.IntegratorState,
              direction: fun_mc.FloatArray,
          ) -> tuple[fun_mc.IntegratorState, fun_mc.IntegratorExtras]:
              directional_step_size = util.map_tree(
                  lambda step_size, state: (  # pylint: disable=g-long-lambda
                      step_size
                      * jnp.reshape(
                          direction,
                          list(direction.shape)
                          + [1]
                          * (len(state.shape) - len(direction.shape)),
                      )
                  ),
                  step_size,
                  state.state,
              )
              # TODO(siege): Ideally we'd pass in the direction here, but the
              # `hamiltonian_integrator` cannot handle dynamic direction switching like
              # that.
              return fun_mc.hamiltonian_integrator(
                  state,
                  num_steps=num_integrator_steps,
                  integrator_step_fn=functools.partial(
                      fun_mc.leapfrog_step,
                      step_size=directional_step_size,
                      target_log_prob_fn=target_log_prob_fn,
                      kinetic_energy_fn=kinetic_energy_fn,
                  ),
                  kinetic_energy_fn=kinetic_energy_fn,
                  unroll=unroll_integrator,
                  max_num_steps=max_num_integrator_steps,
                  integrator_trace_fn=integrator_trace_fn,
              )
  
          integrator_fn = _integrator_fn
  
      seed, sample_seed = util.split_seed(seed, 2)
      momentum = momentum_sample_fn(momentum, sample_seed)
  
      initial_integrator_state = fun_mc.IntegratorState(
          target_log_prob=target_log_prob,
          momentum=momentum,
          state=state,
          state_grads=state_grads,
          state_extra=state_extra,
      )
  
      integrator_state, integrator_extra = integrator_fn(
          initial_integrator_state, direction
      )
  
      proposed_state = phmc_state._replace(
          state=integrator_state.state,
          state_grads=integrator_state.state_grads,
          target_log_prob=integrator_state.target_log_prob,
          momentum=integrator_state.momentum,
          state_extra=integrator_state.state_extra,
          # Flip the direction in the proposal, for reversibility.
          direction=util.map_tree(lambda d: -d, direction),
      )
  
      # Stick the new momentum into phmc_state. We're doing accept/reject purely on
      # the Hamiltonian proposal, not the momentum refreshment kernel.
      phmc_state = phmc_state._replace(momentum=momentum)
  
      energy_change, energy_change_extra = energy_change_fn(
          initial_integrator_state,
          integrator_state,
          integrator_extra,
      )
  
      if log_uniform is None:
          # There's some lint error due to the wrapper here.
          (
              pmh_state,
              pmh_extra,
          ) = fun_mc.persistent_metropolis_hastings_step(  # pytype: disable=wrong-keyword-args
              pmh_state,
              current_state=phmc_state,
              proposed_state=proposed_state,
              energy_change=energy_change,
              drift=mh_drift,
          )
          is_accepted = pmh_extra.is_accepted
          phmc_state = pmh_extra.accepted_state
      else:
          # We explicitly don't update the PMH state.
          phmc_state, mh_extra = fun_mc.metropolis_hastings_step(
              current_state=phmc_state,
              proposed_state=proposed_state,
              energy_change=energy_change,
              log_uniform=log_uniform,
          )
          is_accepted = mh_extra.is_accepted
  
      phmc_state = typing.cast(
          PersistentHamiltonianMonteCarloState, phmc_state
      )
      phmc_state = phmc_state._replace(
          pmh_state=pmh_state,
          # Flip the direction unconditionally; when the state is accepted, this
          # undoes the flip made in the proposal, maintaining the old momentum
          # direction.
          direction=util.map_tree(lambda d: -d, phmc_state.direction),
      )
  
      return phmc_state, PersistentHamiltonianMonteCarloExtra(
          is_accepted=is_accepted,
          proposed_phmc_state=proposed_state,
          log_accept_ratio=-energy_change,
          integrator_state=integrator_state,
          integrator_extra=integrator_extra,
          energy_change_extra=energy_change_extra,
          initial_momentum=momentum,
      )
  
imports: import functools
  import typing
  from typing import Any, Callable, Iterable, NamedTuple, Optional
  import numpy
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import malt
  from fun_mc import sga_hmc
  from fun_mc.malt import *
  from fun_mc.sga_hmc import *
label_desc: Use image node_id 10 for calling a global function with example usage: persistent_hamiltonian_monte_carlo_step(phmc_state, target_log_prob_fn, step_size, num_integrator_steps, noise_fraction, mh_drift, kinetic_energy_fn, momentum_sample_fn, integrator_trace_fn, log_uniform, integrator_fn, unroll_integrator, max_num_integrator_steps, energy_change_fn, named_axis, seed) and returns: phmc_state,  PersistentHamiltonianMonteCarloExtra, new_momentum, fun_mc, unknown