code: def chees_criterion(
      previous_state: fun_mc.State,
      proposed_state: fun_mc.State,
      accept_prob: fun_mc.FloatArray,
      trajectory_length: Optional[fun_mc.FloatArray] = None,
      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, fun_mc.FloatArray]:
      """The ChEES criterion from [1].
  
      ChEES stands for Change in the Estimator of the Expected Square.
  
      ```None
      ChEES = 1/4 E[(||x' - E[x]||**2 - ||x - E[x]||**2)**2],
      ```
  
      where `x` is the previous chain state, `x'` is the next chain state, and
      `||.||` is the L2 norm. Both expectations are with respect to the chain's
      stationary distribution. In practice, the inner expectation is replaced by the
      empirical mean across chains optionally averaged with a provided `state_mean`
      (weighted by `state_mean_weight`).
  
      This can be thought of as the standard expected squared jump distance (ESJD)
      criterion, except that the jump distance is computed in the space of centered
      squared L2 norms. It is also possible to relate ChEES to ESS computed in the
      same space if the true autocorrelation function of the centered squared L2
      norm follows a certain functional form.
  
      ChEES in this implementation is scaled by a normalized
      acceptance probability, so as to discard contributions from bad proposals.
  
      Unlike ChEES, regular ESJD is maximized by perfectly anticorrelated proposals,
      which can give excellent mean estimates but terrible variance estimates;
      maximizing ChEES should give good estimates across a wider range of types of
      posterior expectations.
  
      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: Ignored.
        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 criterion.
        per_chain_chees: The value of the ChEES criterion per chain.
  
      #### References
  
      [1]: Hoffman, M., Radul, A., & Sountsov, P. (2020). An Adaptive MCMC Scheme
           for Setting Trajectory Lengths in Hamiltonian Monte Carlo. In
           preparation.
      """
      del trajectory_length
      batch_ndims = len(accept_prob.shape)
      batch_axes = tuple(range(batch_ndims))
      no_state_mean = object()
      if state_mean is None:
          state_mean = fun_mc.maybe_broadcast_structure(
              no_state_mean, previous_state
          )
      state_mean_weight = fun_mc.maybe_broadcast_structure(
          state_mean_weight, previous_state
      )
      if named_axis is None:
          named_axis_bc = util.map_tree(lambda _: [], previous_state)
      else:
          named_axis_bc = named_axis
  
      if chain_named_axis is None:
          chain_named_axis = []
  
      def _center_previous_state(x, mx, mw):
          x_center = distribute_lib.reduce_mean(
              x, axis=batch_axes, named_axis=chain_named_axis
          )
          if mx is not no_state_mean:
              x_center = x_center * (1 - mw) + mx * mw
          # The empirical mean here is a stand-in for the true mean, so we drop the
          # gradient that flows through this term.
          return x - jax.lax.stop_gradient(x_center)
  
      def _center_proposed_state(x, mx, mw):
          expand_shape = list(accept_prob.shape) + [1] * (
              len(x.shape) - len(accept_prob.shape)
          )
          expanded_accept_prob = jnp.reshape(accept_prob, expand_shape)
  
          # Weight the proposed state by the acceptance probability. The goal here is
          # to get a reliable diagnostic of the underlying dynamics, rather than
          # incorporating the effect of the MetropolisHastings correction.
  
          # accept_prob is zero when x is NaN, but we still want to sanitize such
          # values.
          x_safe = jnp.where(jnp.isfinite(x), x, jnp.zeros_like(x))
          # If all accept_prob's are zero, the x_center will have a nonsense value,
          # but well set the overall criterion to zero in this case, so its fine.
          x_center = distribute_lib.reduce_sum(
              expanded_accept_prob * x_safe,
              axis=batch_axes,
              named_axis=chain_named_axis,
          ) / (
              distribute_lib.reduce_sum(
                  expanded_accept_prob,
                  axis=batch_axes,
                  named_axis=chain_named_axis,
              )
              + 1e-20
          )
          if mx is not no_state_mean:
              x_center = x_center * (1 - mw) + mx * mw
          # The empirical mean here is a stand-in for the true mean, so we drop the
          # gradient that flows through this term.
          return x - jax.lax.stop_gradient(x_center)
  
      def _sum_event_part(x, named_axis):
          event_axes = tuple(range(batch_ndims, len(x.shape)))
          return distribute_lib.reduce_sum(
              x, axis=event_axes, named_axis=named_axis
          )
  
      def _sum_event(x):
          return sum(
              util.flatten_tree(
                  util.map_tree_up_to(
                      x,
                      _sum_event_part,
                      x,
                      named_axis_bc,
                  )
              )
          )
  
      def _square(x):
          return util.map_tree(jnp.square, x)
  
      def _sub(x, y):
          return util.map_tree(lambda x, y: x - y, x, y)
  
      previous_state = util.map_tree(
          _center_previous_state,
          previous_state,
          state_mean,
          state_mean_weight,
      )
      proposed_state = util.map_tree(
          _center_proposed_state,
          proposed_state,
          state_mean,
          state_mean_weight,
      )
      chees = 0.25 * jnp.square(
          _sum_event(
              _sub(_square(proposed_state), _square(previous_state))
          )
      )
  
      # Zero-out per-chain ChEES values where acceptance probability is low. Those
      # values are probably not reflective of the underlying dynamics.
      chees = jnp.where(accept_prob > 1e-4, chees, 0.0)
      accept_prob = accept_prob / distribute_lib.reduce_sum(
          accept_prob + 1e-20, named_axis=chain_named_axis
      )
      chees = chees * accept_prob
  
      return (
          distribute_lib.reduce_mean(
              chees, named_axis=chain_named_axis
          ),
          chees,
      )
  
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 2 for calling a global function with example usage: chees_criterion(previous_state, proposed_state, accept_prob, trajectory_length, state_mean, state_mean_weight, named_axis, chain_named_axis) and returns: distribute_lib,  chees, unknown, unknown, distribute_lib, sum, util, util