code: def snaper_criterion(
      previous_state,
      proposed_state,
      accept_prob,
      trajectory_length,
      direction,
      state_mean=None,
      state_mean_weight=0.0,
      validate_args=False,
      experimental_shard_axis_names=None,
      experimental_reduce_chain_axis_names=None,
  ):
      """The SNAPER criterion from [1].
  
      SNAPER stands for Squared Norm Along Principal component ESJD Rate:
  
      ```None
      SNAPER = E[(((x' - E[x'])^T p)**2 - ((x' - E[x])^T p)**2)**2 /
                 trajectory_length],
      ```
  
      where `x` is the previous chain state, `x'` is the next chain state, and `p`
      is a unit vector (the `direction` argument). 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 unless `state_mean` and
      `state_mean_weight` are set. 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 squared
      projections onto a vector.
  
      The `direction` vector is typically chosen to be an approximation to the first
      principal component of the state covariance matrix.
  
      `state_mean` and `state_mean_weight` can be used to supplement the empirical
      means as follows:
  
      ```None
      E[x] ≈ (1 - state_mean_weight) * x.mean() + state_mean_weight * state_mean.
      ```
  
      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).
        direction: (Possibly nested) floating point `Tensor`. A unit vector onto
          which the centered state should be projected before computing ESJD.
          Typically this chosen to be an approximation to the first principal
          component of the state covariance matrix.
        state_mean: Optional (Possibly nested) floating point `Tensor`. The
          estimated state mean.
        state_mean_weight: Floating point `Tensor`. The weight of the `state_mean`.
        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:
        snaper: The value of the SNAPER criterion.
  
      #### References
  
      [1]: Sountsov, P. & Hoffman, M. (2021). Focusing on Difficult Directions for
           Learning HMC Trajectory Lengths. <https://arxiv.org/abs/2110.11576>
  
      """
      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
      )
  
      if state_mean is None:
          state_mean = tf.nest.map_structure(
              lambda _: None, previous_state
          )
  
          accept_prob = _check_at_least_two_chains(
              accept_prob,
              reduce_chain_axis_names=reduce_chain_axis_names,
              validate_args=validate_args,
              message="snaper_criterion requires at least 2 chains when `state_mean` is `None`",
          )
  
      def _mix_in_state_mean(empirical_mean, state_mean):
          if state_mean is None:
              return empirical_mean
          else:
              return (
                  1.0 - state_mean_weight
              ) * empirical_mean + state_mean_weight * state_mean
  
      def _center_previous_state(x, x_mean):
          # The empirical mean here is a stand-in for the true mean, so we drop the
          # gradient that flows through this term.
          emp_x_mean = tf.stop_gradient(
              distribute_lib.reduce_mean(
                  x, batch_axes, reduce_chain_axis_names
              )
          )
          x_mean = _mix_in_state_mean(emp_x_mean, x_mean)
          return x - x_mean
  
      def _center_proposed_state(x, x_mean):
          # 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.
          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))
          # The empirical mean here is a stand-in for the true mean, so we drop the
          # gradient that flows through this term.
          # 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.
          emp_x_mean = tf.stop_gradient(
              distribute_lib.reduce_sum(
                  expanded_accept_prob * x_safe,
                  batch_axes,
                  reduce_chain_axis_names,
              )
              / (
                  distribute_lib.reduce_sum(
                      expanded_accept_prob,
                      batch_axes,
                      reduce_chain_axis_names,
                  )
                  + 1e-20
              )
          )
  
          x_mean = _mix_in_state_mean(emp_x_mean, x_mean)
          return x - x_mean
  
      def _dot_product_part(x, p, shard_axes=None):
          event_axes = ps.range(batch_ndims, ps.rank(x))
          return distribute_lib.reduce_sum(
              x * p, event_axes, shard_axes
          )
  
      def _dot_product(x):
          dot_products = _map_structure_up_to_with_axes(
              x,
              _dot_product_part,
              x,
              direction,
              experimental_shard_axis_names=experimental_shard_axis_names,
          )
          return sum(tf.nest.flatten(dot_products))
  
      previous_state = tf.nest.map_structure(
          _center_previous_state, previous_state, state_mean
      )
      proposed_state = tf.nest.map_structure(
          _center_proposed_state, proposed_state, state_mean
      )
  
      previous_proj = _dot_product(previous_state)
      proposed_proj = _dot_product(proposed_state)
  
      snaper = (
          tf.square(tf.square(proposed_proj) - tf.square(previous_proj))
          / trajectory_length
      )
      return snaper
  
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 11 for calling a global function with example usage: snaper_criterion(previous_state, proposed_state, accept_prob, trajectory_length, direction, state_mean, state_mean_weight, validate_args, experimental_shard_axis_names, experimental_reduce_chain_axis_names) and returns: snaper, unknown, unknown, distribute_lib, sum, empirical_mean, unknown