code: def stochastic_gradient_ascent_hmc_step(
      sga_hmc_state: StochasticGradientAscentHMCState,
      scalar_step_size: fun_mc.FloatNest,
      criterion_fn: Callable[
          [
              fun_mc.State,
              fun_mc.State,
              fun_mc.FloatArray,
              fun_mc.FloatArray,
          ],
          tuple[fun_mc.FloatArray, Any],
      ],
      trajectory_length_adaptation_rate: float
      | fun_mc.FloatArray = 0.05,
      trajectory_length_sample_fn: Callable[
          [Any, fun_mc.IntArray, Any], fun_mc.FloatArray
      ] = (default_trajectory_length_sample),
      trajectory_length_constrain_fn: Callable[[Any], Any] = (
          default_trajectory_length_constrain
      ),
      adam_kwargs: Mapping[str, Any] = immutabledict.immutabledict(
          {"beta_1": 0.0, "beta_2": 0.5}
      ),
      averaging_window_steps: int | fun_mc.IntArray = 100,
      adapt: bool | fun_mc.BooleanArray = True,
      seed: Any = None,
      **hmc_kwargs: Mapping[str, Any],
  ):
      """Stochastic gradient ascent Hamiltonian Monte Carlo step.
  
      SGA-HMC posits an existence of a parameteric distribution over trajectory
      lengths. It then uses stochastic gradients to adapt those parameters by
      maximizing the expected value of some criterion.
  
      The gradients are computed by the use of
      `hamiltonian_monte_carlo_with_state_grads_step` and then using Monte-Carlo
      averages across separate Markov Chains and Markov Chain iterates.
  
      ChEES [1] criterion is the prototypical example.
  
      The trajectory distribution is parameterized via `trajectory_length_sample_fn`
      and `trajectory_length_params_constrain_fn`. While `adapt` is `False`, the
      parameters are adapted using Adam (controlled using
      `trajectory_length_adaptation_rate` and `adam_kwargs`). When `adapt` is
      `False`, averaged parameters are used, which have been computed via an
      exponential moving while `adapt` was `True`. The degree of averaging is
      controlled via `averaging_window_steps`. The parameters that were actually
      used for this step are returned in the `trajectory_length_params` field in the
      `sga_hmc_extra` return.
  
      Args:
        sga_hmc_state: `StochasticGradientAscentHMCState`
        scalar_step_size: Scalar step size (see
          `hamiltonian_monte_carlo_with_state_grads_step` for details).
        criterion_fn: Callable with signature `(previous_state, proposed_state,
          accept_prob, trajectory_length) -> (criterion, criterion_extra)`. The
          criterion to maximize.
        trajectory_length_adaptation_rate: Adaption rate for the trajectory length
          parameters.
        trajectory_length_sample_fn: Callable with signature
          `(trajectory_length_params, step, seed) -> trajectory_length`. Used to
          sample a new trajectory length.
        trajectory_length_constrain_fn: Used to constrain the trajectory length
          parameters by projecting them into the allowed set.
        adam_kwargs: Additional keyword arguments for Adam optimizer used to adapt
          the trajectory length parameters.
        averaging_window_steps: Window size for averaging the trajectory parameters.
          See `fun_mc.running_mean_step` for the meaning of this argument.
        adapt: Whether to adapt the trajectory parameters and whether to use the
          adapted parameters or the averaged parameters when sampling the trajectory
          for this step.
        seed: PRNG seed.
        **hmc_kwargs: Passed to `hamiltonian_monte_carlo_with_state_grads_step`.
  
      Returns:
        sga_hmc_state: `StochasticGradientAscentHMCState`.
        sga_hmc_extra: `StochasticGradientAscentHMCExtra`.
  
      #### References
  
      [1]: Hoffman, M., Radul, A., & Sountsov, P. (2020). An Adaptive MCMC Scheme
           for Setting Trajectory Lengths in Hamiltonian Monte Carlo. In
           preparation.
      """
  
      seed, sample_seed, hmc_seed = util.split_seed(seed, 3)
  
      @util.named_call
      def loss_fn(*args, **kwargs):
          rmean_params = (
              sga_hmc_state.trajectory_length_params_rmean_state.mean
          )
          adapting_params = fun_mc.recover_state_from_args(
              args, kwargs, rmean_params
          )
          params = fun_mc.choose(adapt, adapting_params, rmean_params)
          trajectory_length = trajectory_length_sample_fn(
              params, sga_hmc_state.step, sample_seed
          )
  
          (
              hmc_state,
              hmc_extra,
          ) = hamiltonian_monte_carlo_with_state_grads_step(  # pytype: disable=wrong-keyword-args
              sga_hmc_state.hmc_state,
              trajectory_length=trajectory_length,
              scalar_step_size=scalar_step_size,
              seed=hmc_seed,
              **hmc_kwargs,
          )
  
          accept_prob = jnp.exp(
              jnp.minimum(
                  jnp.zeros_like(hmc_extra.hmc_extra.log_accept_ratio),
                  hmc_extra.hmc_extra.log_accept_ratio,
              )
          )
          accept_prob = jnp.where(
              jnp.isfinite(accept_prob),
              accept_prob,
              jnp.zeros_like(accept_prob),
          )
  
          criterion, criterion_extra = criterion_fn(
              sga_hmc_state.hmc_state.state,
              hmc_extra.proposed_state,
              accept_prob,
              # + step_size because we're effectively doing floor(traj / step_size)
              # when computing the number of leapfrog steps.
              trajectory_length + scalar_step_size,
          )
  
          return -criterion, (
              hmc_state,
              hmc_extra,
              criterion,
              criterion_extra,
              params,
          )
  
      # Adapt trajectory.
      (
          trajectory_length_params_opt_state,
          trajectory_length_params_opt_extra,
      ) = fun_mc.adam_step(
          sga_hmc_state.trajectory_length_params_opt_state,
          loss_fn,
          learning_rate=trajectory_length_adaptation_rate,
          **adam_kwargs,
      )
  
      (
          hmc_state,
          hmc_extra,
          criterion,
          criterion_extra,
          trajectory_length_params,
      ) = trajectory_length_params_opt_extra.loss_extra
  
      # Constrain trajectory params.
      trajectory_length_params_opt_state = fun_mc.choose(
          adapt,
          trajectory_length_params_opt_state,
          sga_hmc_state.trajectory_length_params_opt_state,
      )
      constrained_trajectory_length_params = (
          trajectory_length_constrain_fn(
              trajectory_length_params_opt_state.state
          )
      )
      trajectory_length_params_opt_state = (
          trajectory_length_params_opt_state._replace(
              state=constrained_trajectory_length_params
          )
      )
  
      # Update the running mean for trajectory params.
      (
          trajectory_length_params_rmean_state,
          _,
      ) = fun_mc.running_mean_step(
          sga_hmc_state.trajectory_length_params_rmean_state,
          trajectory_length_params_opt_state.state,
          window_size=averaging_window_steps,
      )
      trajectory_length_params_rmean_state = fun_mc.choose(
          adapt,
          trajectory_length_params_rmean_state,
          sga_hmc_state.trajectory_length_params_rmean_state,
      )
  
      sga_hmc_state = sga_hmc_state._replace(
          hmc_state=hmc_state,
          step=sga_hmc_state.step + 1,
          trajectory_length_params_rmean_state=trajectory_length_params_rmean_state,
          trajectory_length_params_opt_state=trajectory_length_params_opt_state,
      )
      sga_hmc_extra = StochasticGradientAscentHMCExtra(
          hmc_extra=hmc_extra.hmc_extra,
          num_integrator_steps=hmc_extra.num_integrator_steps,
          trajectory_length_params_opt_extra=trajectory_length_params_opt_extra,
          criterion=criterion,
          criterion_extra=criterion_extra,
          trajectory_length_params=trajectory_length_params,
      )
      return sga_hmc_state, sga_hmc_extra
  
imports: from collections.abc import Mapping
  from typing import Any, Callable, NamedTuple, Optional
  import immutabledict
  from fun_mc import backend
  from fun_mc import fun_mc_lib
label_desc: Use image node_id 9 for calling a global function with example usage: stochastic_gradient_ascent_hmc_step(sga_hmc_state, scalar_step_size, criterion_fn, trajectory_length_adaptation_rate, trajectory_length_sample_fn, trajectory_length_constrain_fn, adam_kwargs, averaging_window_steps, adapt, seed) and returns: sga_hmc_state,  sga_hmc_extra, unknown,  hmc_state,  hmc_extra,  criterion,  criterion_extra,  params