code: def hamiltonian_monte_carlo_with_state_grads_step(
      hmc_state: fun_mc.HamiltonianMonteCarloState,
      trajectory_length: fun_mc.FloatArray,
      scalar_step_size: fun_mc.FloatArray,
      step_size_scale: float | fun_mc.FloatNest = 1.0,
      named_axis: Optional[fun_mc.StringNest] = None,
      **hmc_kwargs,
  ) -> tuple[
      fun_mc.HamiltonianMonteCarloState,
      HamiltonianMonteCarloWithStateGradsExtra,
  ]:
      """Hamiltonian Monte Carlo (HMC) step with gradients for proposed state.
  
      This acts as a `fun_mc.hamiltonian_monte_carlo_step`, where the
      `num_integrator_steps` is defined as `ceil(trajectory_length /
      scalar_step_size)` and `step_size` is defined as `scalar_step_size *
      step_size_scale`. The main feature of this function is that it propagates the
      gradients from `hmc_with_state_grads_extra.proposed_state` to
      `trajectory_length` (these are the only gradients propagated at the moment).
      This feature can be used to do gradient-based optimization of
      `trajectory_length` based on criteria that depend on the `proposed_state`
      (e.g. [1]).
  
      Args:
        hmc_state: `fun_mc.HamiltonianMonteCarloState`.
        trajectory_length: Trajectory length used by HMC.
        scalar_step_size: Scalar step size (used to compute the number of leapfrog
          steps).
        step_size_scale: Step size scale, structure broadcastable to the
          `hmc_state.state`.
        named_axis: Named axes of the state. Same structure as `hmc_state.state`.
        **hmc_kwargs: Passed to `fun_mc.hamiltonian_monte_carlo_step`.
  
      Returns:
        hmc_state: `fun_mc.HamiltonianMonteCarloState`.
        hmc_with_grads_extra: Extra outputs.
  
      #### References
  
      [1]: Hoffman, M., Radul, A., & Sountsov, P. (2021). An Adaptive MCMC Scheme
           for Setting Trajectory Lengths in Hamiltonian Monte Carlo.
           http://proceedings.mlr.press/v130/hoffman21a.html
      """
      consts = scalar_step_size, hmc_state, step_size_scale
      flat_consts = util.flatten_tree(consts)
  
      @jax.custom_gradient
      def hmc(*traj_and_flat_consts):
          trajectory_length = traj_and_flat_consts[0]
          (
              scalar_step_size,
              hmc_state,
              step_size_scale,
          ) = util.unflatten_tree(consts, traj_and_flat_consts[1:])
          trajectory_length = jnp.asarray(trajectory_length)
          num_integrator_steps = jnp.asarray(
              jnp.ceil(trajectory_length / scalar_step_size), jnp.int32
          )
          # In case something goes negative.
          num_integrator_steps = jnp.maximum(1, num_integrator_steps)
          (
              new_hmc_state,
              hmc_extra,
          ) = fun_mc.hamiltonian_monte_carlo_step(
              hmc_state,
              num_integrator_steps=num_integrator_steps,
              step_size=util.map_tree(
                  lambda s: s * scalar_step_size, step_size_scale
              ),
              named_axis=named_axis,
              **hmc_kwargs,
          )
          hmc_with_grads_extra = (
              HamiltonianMonteCarloWithStateGradsExtra(
                  proposed_state=hmc_extra.proposed_hmc_state.state,
                  hmc_extra=hmc_extra,
                  num_integrator_steps=num_integrator_steps,
              )
          )
          res = (new_hmc_state, hmc_with_grads_extra)
  
          def grad(*grads):
              grads = util.unflatten_tree(res, util.flatten_tree(grads))
  
              step_size_scale_bc = fun_mc.maybe_broadcast_structure(
                  step_size_scale,
                  hmc_extra.integrator_extra.momentum_grads,
              )
  
              # We wish to compute `grads^T @
              # jacobian(proposed_state(trajectory_length))`.
              #
              # The Jacobian is known from from Hamilton's equations:
              #
              # dx / dt = dK(v) / dv
              #
              # where `x` is the state, `v` is the momentum and `K` is the kinetic
              # energy. Since `step_size_scale` rescales momentum, we the right hand
              # side of that expression is `momentum_grads * step_size_scale` by the
              # chain rule. Since the Jacobian in question has 1 row, the
              # vector-Jacobian product is simply the dot product.
              state_grads = util.map_tree(
                  lambda s, m, g: s * m * g,
                  step_size_scale_bc,
                  hmc_extra.integrator_extra.momentum_grads,
                  grads[1].proposed_state,
              )
  
              def do_sum(x, named_axis):
                  return distribute_lib.reduce_sum(
                      x,
                      list(
                          range(
                              len(trajectory_length.shape), len(x.shape)
                          )
                      ),
                      named_axis,
                  )
  
              if named_axis is None:
                  named_axis_bc = util.map_tree(
                      lambda _: [], state_grads
                  )
              else:
                  named_axis_bc = named_axis
  
              traj_grad = sum(
                  util.flatten_tree(
                      util.map_tree_up_to(
                          state_grads,
                          do_sum,
                          state_grads,
                          named_axis_bc,
                      )
                  )
              )
              return (traj_grad,) + (None,) * len(flat_consts)
  
          return res, grad
  
      return hmc(trajectory_length, *flat_consts)
  
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 1 for calling a global function with example usage: hamiltonian_monte_carlo_with_state_grads_step(hmc_state, trajectory_length, scalar_step_size, step_size_scale, named_axis) and returns: hmc, res,  grad, unknown, distribute_lib