code: def snaper_criterion(  # pytype: disable=annotation-type-mismatch  # jax-ndarray
      previous_state: jnp.ndarray,
      proposed_state: jnp.ndarray,
      accept_prob: jnp.ndarray,
      trajectory_length: jnp.ndarray,
      principal: jnp.ndarray,
      power: jnp.ndarray = 1.0,
      state_mean: jnp.ndarray = Optional[jnp.ndarray],
      state_mean_weight: jnp.ndarray = 0.0,
  ):
      """SNAPER criterion[1].
  
      Args:
        previous_state: Previous MCMC state. Shape: [B, D]
        proposed_state: Proposed MCMC state. Shape: [B, D]
        accept_prob: Acceptance probability. Shape: [B]
        trajectory_length: Trajectory length. Shape: []
        principal: Principal component. Shape: [D]
        power: Power to raise the trajectory length to when normalizing. Shape: []
        state_mean: Optional mean of the state. Computed across the B dimension
          otherwise.
        state_mean_weight: Weight in [0, 1] to use when mixing `state_mean` and the
          empirical average across the B dimension.
  
      Returns:
        A tuple of (criterion, ()).
  
      #### References:
  
      [1]: Sountsov, P., & Hoffman, M. D. (2021). Focusing on Difficult Directions
           for Learning HMC Trajectory Lengths. In arXiv [stat.CO]. arXiv.
           http://arxiv.org/abs/2110.11576
      """
  
      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
      )
      mx = state_mean
      mw = state_mean_weight
  
      x_center = previous_state.mean(0)
      if mx is not no_state_mean:
          x_center = x_center * (1 - mw) + mx * mw
      previous_state = previous_state - x_center
  
      x = proposed_state
      expand_shape = list(accept_prob.shape) + [1] * (
          len(x.shape) - len(accept_prob.shape)
      )
      expanded_accept_prob = jnp.reshape(accept_prob, expand_shape)
  
      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 we'll set the overall criterion to zero in this case, so it's fine.
      x_center = jnp.sum(
          expanded_accept_prob * x_safe,
          axis=batch_axes,
      ) / (jnp.sum(expanded_accept_prob, axis=batch_axes) + 1e-20)
      if mx is not no_state_mean:
          x_center = x_center * (1 - mw) + mx * mw
      proposed_state = x - jax.lax.stop_gradient(x_center)
  
      previous_projection = jnp.einsum(
          "d,nd->n", principal, previous_state
      )
      proposed_projection = jnp.einsum(
          "d,nd->n", principal, proposed_state
      )
  
      esjd = (previous_projection**2 - proposed_projection**2) ** 2
  
      esjd = jnp.where(accept_prob > 1e-4, esjd, 0.0)
      accept_prob = accept_prob / jnp.sum(accept_prob + 1e-20)
      esjd = esjd * accept_prob
  
      return esjd.mean() / trajectory_length**power, {
          "previous_projection": previous_projection,
          "proposed_projection": proposed_projection,
      }
  
imports: import functools
  import os
  from typing import Any, Callable, Dict, Optional, Mapping, NamedTuple, Tuple, Union
  from etils import epath
  import gin
  import immutabledict
  import jax
  import jax.numpy
  import numpy
  from discussion.adaptive_malt import utils
  from fun_mc import using_jax
  from inference_gym import using_jax
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 4 for calling a global function with example usage: snaper_criterion(previous_state, proposed_state, accept_prob, trajectory_length, principal, power, state_mean, state_mean_weight) and returns: unknown,  dict