code: def chees_per_grad_criterion(
      previous_state: fun_mc.State,
      proposed_state: fun_mc.State,
      accept_prob: fun_mc.FloatArray,
      trajectory_length: fun_mc.FloatArray,
      power: float | fun_mc.FloatArray = 1.0,
      state_mean: Optional[fun_mc.State] = None,
      state_mean_weight: float | fun_mc.FloatNest = 0.0,
      named_axis: Optional[fun_mc.StringNest] = None,
      chain_named_axis: Optional[fun_mc.StringNest] = None,
  ) -> tuple[fun_mc.FloatArray, ChEESPerGradExtra]:
      """ChEES per gradient criterion.
  
      This criterion is computed as:
  
      ```none
      ChEES/grad = 1/4 E[(||x' - E[x]||**2 - ||x - E[x]||**2)**2 /
        trajectory_length**power].
      ```
  
      For constant step sizes, `trajectory_length` is proportional to the number of
      gradients evaluated, meaning that this criterion trades off maximizing ChEES
      and using more computation to do so. In practice, this mitigates an issue with
      ChEES where it is periodic and has spurious maxima at prohibitively large
      trajectories that can nonetheless be found by gradient ascent with an
      aggressive learning rate. This criterion penalizes those maxima, making them
      less likely to be found.
  
      Args:
        previous_state: (Possibly nested) floating point `Array`. The previous state
          of the MCMC chain.
        proposed_state: (Possibly nested) floating point `Array`. The proposed state
          of the MCMC chain.
        accept_prob: Floating `Array`. Probability of acceping the proposed state.
        trajectory_length: Trajectory length associated with the transition from
          `previous_state` to `proposed_state`.
        power: Floating `Array`. Used to scale the `trajectory_length` term.
        state_mean: (Possibly nested) floating point `Array`. Optional estimate of
          the MCMC chain mean.
        state_mean_weight: Floating point `Array`. Used to weight `state_mean` with
          the mean computed by averaging across the previous/proposed state. Setting
          it to effectively uses `state_mean` as the only source of the MCMCM chain
          mean.
        named_axis: Named axes of the state. Same structure as `previous_state`.
        chain_named_axis: Named axes of the MCMC chain that the criterion is to be
          averaged over.
  
      Returns:
        chees: The value of the ChEES per grad criterion.
        chees_per_grad_extra: `ChEESPerGradExtra`.
      """
      chees, per_chain_chees = chees_criterion(
          previous_state=previous_state,
          proposed_state=proposed_state,
          accept_prob=accept_prob,
          state_mean=state_mean,
          state_mean_weight=state_mean_weight,
          named_axis=named_axis,
          chain_named_axis=chain_named_axis,
      )
      per_chain_chees_per_grad = (
          per_chain_chees
          / distribute_lib.pbroadcast(
              trajectory_length**power, chain_named_axis
          )
      )
      extra = ChEESPerGradExtra(
          chees=chees,
          per_chain_chees=per_chain_chees,
          per_chain_chees_per_grad=per_chain_chees_per_grad,
      )
      return (
          distribute_lib.reduce_mean(
              per_chain_chees_per_grad, named_axis=chain_named_axis
          ),
          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 3 for calling a global function with example usage: chees_per_grad_criterion(previous_state, proposed_state, accept_prob, trajectory_length, power, state_mean, state_mean_weight, named_axis, chain_named_axis) and returns: distribute_lib,  extra