code: def chees_criterion(
      previous_state,
      proposed_state,
      accept_prob,
      trajectory_length,
      validate_args=False,
      experimental_shard_axis_names=None,
      experimental_reduce_chain_axis_names=None,
  ):
      """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, so computing this criterion requires that at
      least 2 chains are present. The outer expectation is computed by the caller
      (e.g. in the `GradientBasedTrajectoryLengthAdaptation` kernel).
  
      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.
  
      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 `Tensor`. The previous
          state of the HMC chain.
        proposed_state: (Possibly nested) floating point `Tensor`. The proposed
          state of the HMC chain.
        accept_prob: Floating `Tensor`. Probability of acceping the proposed state.
        trajectory_length: Floating `Tensor`. Mean trajectory length (not used in
          this criterion).
        validate_args: Whether to perform non-static argument validation.
        experimental_shard_axis_names: A structure of string names indicating how
          members of the state are sharded.
        experimental_reduce_chain_axis_names: A string or list of string names
          indicating which named chain axes to reduce over when computing the
          criterion.
  
      Returns:
        chees: The value of the ChEES criterion.
  
      Raises:
        ValueError: If `accept_prob` indicates that there are fewer than 2 chains.
  
      #### References
  
      [1]: Hoffman, M., Radul, A., & Sountsov, P. (2020). An Adaptive MCMC Scheme
           for Setting Trajectory Lengths in Hamiltonian Monte Carlo.
           <https://proceedings.mlr.press/v130/hoffman21a>
  
      """
      del trajectory_length
      batch_ndims = ps.rank(accept_prob)
      batch_axes = ps.range(batch_ndims, dtype=tf.int32)
      reduce_chain_axis_names = distribute_lib.canonicalize_named_axis(
          experimental_reduce_chain_axis_names
      )
  
      accept_prob = _check_at_least_two_chains(
          accept_prob,
          reduce_chain_axis_names=reduce_chain_axis_names,
          validate_args=validate_args,
          message="chees_criterion requires at least 2 chains.",
      )
  
      def _center_previous_state(x):
          # The empirical mean here is a stand-in for the true mean, so we drop the
          # gradient that flows through this term.
          x_mean = _reduce_mean_with_axes(
              x, batch_axes, reduce_chain_axis_names
          )
          return x - tf.stop_gradient(x_mean)
  
      def _center_proposed_state(x):
          # Note that we don't do a monte carlo average of the accepted chain
          # position, but rather try to get an estimate of the underlying dynamics.
          # This is done by only looking at proposed states where the integration
          # error is low.
          # TODO(mhoffman): Needs more experimentation.
          expanded_accept_prob = bu.left_justified_expand_dims_like(
              accept_prob, x
          )
  
          # accept_prob is zero when x is NaN, but we still want to sanitize such
          # values.
          x_safe = tf.where(tf.math.is_finite(x), x, tf.zeros_like(x))
          # If all accept_prob's are zero, the x_center will have a nonsense value,
          # but we'll discard the resultant gradients later on, so it's fine.
          x_center = _reduce_sum_with_axes(
              expanded_accept_prob * x_safe,
              batch_axes,
              reduce_chain_axis_names,
          ) / (
              _reduce_sum_with_axes(
                  expanded_accept_prob,
                  batch_axes,
                  reduce_chain_axis_names,
              )
              + 1e-20
          )
  
          # The empirical mean here is a stand-in for the true mean, so we drop the
          # gradient that flows through this term.
          return x - tf.stop_gradient(x_center)
  
      def _sum_event_part(x, shard_axes=None):
          event_axes = ps.range(batch_ndims, ps.rank(x))
          return distribute_lib.psum(
              tf.reduce_sum(x, axis=event_axes), shard_axes
          )
  
      def _sum_event(x):
          event_parts = _map_structure_up_to_with_axes(
              x,
              _sum_event_part,
              x,
              experimental_shard_axis_names=experimental_shard_axis_names,
          )
          return sum(tf.nest.flatten(event_parts))
  
      def _square(x):
          return tf.nest.map_structure(tf.square, x)
  
      def _sub(x, y):
          return tf.nest.map_structure(lambda x, y: x - y, x, y)
  
      previous_state = tf.nest.map_structure(
          _center_previous_state, previous_state
      )
      proposed_state = tf.nest.map_structure(
          _center_proposed_state, proposed_state
      )
      chees = 0.25 * tf.square(
          _sum_event(
              _sub(_square(proposed_state), _square(previous_state))
          )
      )
      return chees
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 9 for calling a global function with example usage: chees_criterion(previous_state, proposed_state, accept_prob, trajectory_length, validate_args, experimental_shard_axis_names, experimental_reduce_chain_axis_names) and returns: chees, unknown, unknown, distribute_lib, sum, tf, tf