code: def adaptive_hamiltonian_monte_carlo_step(
      adaptive_hmc_state: AdaptiveHamiltonianMonteCarloState,
      target_log_prob_fn: fun_mc.PotentialFn,
      num_adaptation_steps: Optional[fun_mc.IntArray],
      variance_window_steps: int | fun_mc.IntArray = 100,
      trajectory_length_factor: float | fun_mc.FloatArray = 1.0,
      num_trajectory_ramp_steps: Optional[int] = 100,
      trajectory_warmup_power: float | fun_mc.FloatArray = 1.0,
      max_num_leapfrog_steps: Optional[int] = 100,
      step_size_adaptation_rate: float | fun_mc.FloatArray = 0.05,
      step_size_adaptation_rate_decay_power: float
      | fun_mc.FloatArray = 0.1,
      target_accept_prob: float | fun_mc.FloatArray = 0.8,
      step_size_averaging_window_steps: int | fun_mc.IntArray = 100,
      jitter_sample_fn: Callable[
          [fun_mc.IntArray, fun_mc.IntArray, Any], fun_mc.IntArray
      ] = (_uniform_jitter),
      log_accept_ratio_reduce_fn: Callable[
          [fun_mc.FloatArray], fun_mc.FloatArray
      ] = (tfp.math.reduce_logmeanexp),
      hmc_kwargs: Optional[dict[str, Any]] = None,
      seed: Any = None,
  ) -> tuple[
      AdaptiveHamiltonianMonteCarloState,
      AdaptiveHamiltonianMonteCarloExtra,
  ]:
      """Adaptive Hamiltonian Monte Carlo `TransitionOperator`.
  
      This implements a relatively straighforward adaptive HMC algorithm with
      diagonal mass-matrix adaptation and step size adaptation. The algorithm also
      estimates the trajectory length based on the variance of the chain.
  
      All adaptation stops after `num_adaptation_steps`, after which this algorithm
      becomes regular HMC with fixed number of leapfrog steps, and fixed
      per-component step size. Typically, chain samples after `num_adaptation_steps
      + num_warmup_steps` are discarded, where `num_warmup_steps` is typically
      heuristically chosen to be `0.25 * num_adaptation_steps`. For maximum
      efficiency, however, it's recommended to actually use
      `fun_mc.hamiltonian_monte_carlo_step` initialized with the relevant
      hyperparameters, as that `TransitionOperator` won't have the overhead of the
      adaptation logic. This can be set to `None`, in which case adaptation never
      stops (and the algorihthm ceases to be calibrated).
  
      The mass matrix is adapted by computing an exponential moving variance of the
      chain. The averaging window is controlled by the `variance_window_steps`, with
      larger values leading to a smoother estimate.
  
      Trajectory length is computed as `max(sqrt(chain_variance)) *
      trajectory_length_factor`. To be resilient to poor initialization,
      `trajectory_length_factor` can be increased from 0 based on a polynomial
      schedule, controlled by `num_trajectory_ramp_steps` and
      `trajectory_warmup_power`.
  
      The step size is adapted to make the acceptance probability close to
      `target_accept_prob`, using the `Adam` optimizer. This is controlled by the
      `step_size_adaptation_rate`. If `num_adaptation_steps` is not `None`, this
      rate is decayed using a polynomial schedule controlled by
      `step_size_adaptation_rate`.
  
      Args:
        adaptive_hmc_state: `AdaptiveHamiltonianMonteCarloState`
        target_log_prob_fn: Target log prob fn.
        num_adaptation_steps: Number of adaptation steps, can be `None`.
        variance_window_steps: Window to compute the chain variance over.
        trajectory_length_factor: Trajectory length factor.
        num_trajectory_ramp_steps: Number of steps to warmup the
          `trajectory_length_factor`.
        trajectory_warmup_power: Power of the polynomial schedule for
          `trajectory_length_factor` warmup.
        max_num_leapfrog_steps: Maximum number of leapfrog steps to take.
        step_size_adaptation_rate: Step size adaptation rate.
        step_size_adaptation_rate_decay_power:  Power of the polynomial schedule for
          `trajectory_length_factor` warmup.
        target_accept_prob: Target acceptance probability.
        step_size_averaging_window_steps: Number of steps to compute the averaged
          step size.
        jitter_sample_fn: Function with signature `(mean_num_leapfrog_steps, step,
          seed) -> num_leapfrog_steps`. By default, this does a uniform jitter
          returning a step size in [1, 2 * mean_num_leapfrog_steps).
        log_accept_ratio_reduce_fn: A function that reduces `log_accept_ratio` in
          log-space. By default, this computes the log-mean-exp.
        hmc_kwargs: Additional keyword arguments to pass to
          `hamiltonian_monte_carlo_step`.
        seed: Random seed to use.
  
      Returns:
        adaptive_hmc_state: `AdaptiveHamiltonianMonteCarloState`.
        adaptive_hmc_extra: `AdaptiveHamiltonianMonteCarloExtra`.
  
      #### Examples
  
      Here's an example using using Adaptive HMC and TensorFlow Probability to
      sample from a simple model.
  
      ```python
      num_chains = 16
      num_steps = 2000
      num_warmup_steps = num_steps // 2
      num_adapt_steps = int(0.8 * num_warmup_steps)
  
      # Setup the model and state constraints.
      model = tfp.distributions.JointDistributionSequential([
          tfp.distributions.Normal(loc=0., scale=1.),
          tfp.distributions.Independent(
              tfp.distributions.LogNormal(loc=[1., 1.], scale=0.5), 1),
      ])
      bijector = [tfp.bijectors.Identity(), tfp.bijectors.Exp()]
      transform_fn = fun_mc.util_tfp.bijector_to_transform_fn(
          bijector, model.dtype, batch_ndims=1)
  
      def target_log_prob_fn(*x):
        return model.log_prob(x), ()
  
      # Start out at zeros (in the unconstrained space).
      state, _ = transform_fn(
          *map(lambda e: jnp.zeros([num_chains] + list(e)), model.event_shape))
  
      reparam_log_prob_fn, reparam_state = fun_mc.reparameterize_potential_fn(
          target_log_prob_fn, transform_fn, state)
  
      # Define the kernel.
      def kernel(adaptive_hmc_state):
        adaptive_hmc_state, adaptive_hmc_extra = (
            fun_mc.prefab.adaptive_hamiltonian_monte_carlo_step(
                adaptive_hmc_state,
                target_log_prob_fn=reparam_log_prob_fn,
                num_adaptation_steps=num_adapt_steps))
  
        return adaptive_hmc_state, (adaptive_hmc_extra.state,
                                    adaptive_hmc_extra.is_accepted,
                                    adaptive_hmc_extra.step_size)
  
  
      _, (state_chain, is_accepted_chain, step_size_chain) = jax.jit(
          lambda: fun_mc.trace(
              state=fun_mc.prefab.adaptive_hamiltonian_monte_carlo_init(
                  reparam_state, reparam_log_prob_fn),
              fn=kernel,
              num_steps=num_steps),
          autograph=False)()
  
      # Discard the warmup samples.
      state_chain = [s[num_warmup_steps:] for s in state_chain]
      is_accepted_chain = is_accepted_chain[num_warmup_steps:]
  
      # Compute diagnostics.
      accept_rate = jnp.mean(jnp.array(is_accepted_chain, jnp.float32))
      ess = tfp.mcmc.effective_sample_size(
          state_chain, filter_beyond_positive_pairs=True, cross_chain_dims=[1, 1])
      rhat = tfp.mcmc.potential_scale_reduction(state_chain)
  
      # Compute relevant quantities.
      sample_mean = [jnp.mean(s, axis=[0, 1]) for s in state_chain]
      sample_var = [jnp.var(s, axis=[0, 1]) for s in state_chain]
  
      # It's also important to look at the `step_size_chain` (e.g. via a plot), to
      # verify that adaptation succeeded.
      ```
      """
      dtype = util.flatten_tree(adaptive_hmc_state.hmc_state.state)[
          0
      ].dtype
      trajectory_length_factor = jnp.asarray(
          trajectory_length_factor, dtype=dtype
      )
      trajectory_warmup_power = jnp.asarray(
          trajectory_warmup_power, dtype=dtype
      )
  
      hmc_state = adaptive_hmc_state.hmc_state
      running_var_state = adaptive_hmc_state.running_var_state
      ssa_state = adaptive_hmc_state.ssa_state
      step = adaptive_hmc_state.step
  
      # Warmup the trajectory length, if requested.
      if num_trajectory_ramp_steps is not None:
          trajectory_length_factor = _polynomial_decay(
              step=step,
              step_size=jnp.array(0.0, dtype),
              decay_steps=num_trajectory_ramp_steps,
              final_step_size=trajectory_length_factor,
              power=trajectory_warmup_power,
          )
  
      # Compute the per-component step_size and num_leapfrog_steps from the variance
      # estimate.
      scale = util.map_tree(jnp.sqrt, running_var_state.variance)
      step_size = ssa_state.step_size(
          num_adaptation_steps=num_adaptation_steps
      )
      num_leapfrog_steps = jnp.array(
          jnp.ceil(trajectory_length_factor / step_size), jnp.int32
      )
      num_leapfrog_steps = jnp.maximum(1, num_leapfrog_steps)
      if max_num_leapfrog_steps is not None:
          num_leapfrog_steps = jnp.minimum(
              max_num_leapfrog_steps, num_leapfrog_steps
          )
      # We implement mass-matrix adaptation via step size rescaling, as this is a
      # little bit simpler to code up.
      step_size = util.map_tree(lambda scale: scale * step_size, scale)
  
      hmc_seed, jitter_seed = util.split_seed(seed, 2)
      jittered_num_leapfrog_steps = jitter_sample_fn(
          num_leapfrog_steps, step, jitter_seed
      )
  
      # Run a step of HMC.
      hmc_kwargs = hmc_kwargs or {}
      hmc_state, hmc_extra = fun_mc.hamiltonian_monte_carlo_step(
          hmc_state,
          target_log_prob_fn=target_log_prob_fn,
          step_size=step_size,
          num_integrator_steps=jittered_num_leapfrog_steps,
          seed=hmc_seed,
          **hmc_kwargs,
      )
  
      # Update the running variance estimate.
      chain_ndims = len(hmc_state.target_log_prob.shape)
      old_running_var_state = running_var_state
      running_var_state, _ = fun_mc.running_variance_step(
          running_var_state,
          hmc_state.state,
          axis=tuple(range(chain_ndims)) if chain_ndims else None,
          window_size=int(np.prod(hmc_state.target_log_prob.shape))
          * variance_window_steps,
      )
  
      if num_adaptation_steps is not None:
          # Take care of adaptation for variance.
          running_var_state = util.map_tree(
              lambda n, o: jnp.where(
                  step < num_adaptation_steps, n, o
              ),  # pylint: disable=g-long-lambda
              running_var_state,
              old_running_var_state,
          )
  
      # Update the scalar step size as a function of acceptance rate.
      ssa_state, _ = step_size_adaptation_step(
          ssa_state,
          log_accept_ratio=hmc_extra.log_accept_ratio,
          num_adaptation_steps=num_adaptation_steps,
          target_accept_prob=target_accept_prob,
          adaptation_rate=step_size_adaptation_rate,
          adaptation_rate_decay_power=step_size_adaptation_rate_decay_power,
          reduce_fn=log_accept_ratio_reduce_fn,
          averaging_window_steps=step_size_averaging_window_steps,
      )
  
      adaptive_hmc_state = AdaptiveHamiltonianMonteCarloState(
          hmc_state=hmc_state,
          running_var_state=running_var_state,
          ssa_state=ssa_state,
          step=step + 1,
      )
  
      extra = AdaptiveHamiltonianMonteCarloExtra(
          hmc_state=hmc_state,
          hmc_extra=hmc_extra,
          step_size=step_size,
          mean_num_leapfrog_steps=num_leapfrog_steps,
          num_leapfrog_steps=jittered_num_leapfrog_steps,
      )
  
      return adaptive_hmc_state, extra
  
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 6 for calling a global function with example usage: adaptive_hamiltonian_monte_carlo_step(adaptive_hmc_state, target_log_prob_fn, num_adaptation_steps, variance_window_steps, trajectory_length_factor, num_trajectory_ramp_steps, trajectory_warmup_power, max_num_leapfrog_steps, step_size_adaptation_rate, step_size_adaptation_rate_decay_power, target_accept_prob, step_size_averaging_window_steps, jitter_sample_fn, log_accept_ratio_reduce_fn, hmc_kwargs, seed) and returns: adaptive_hmc_state,  extra