code: def _compute_log_acceptance_correction(
      current_momentums,
      proposed_momentums,
      independent_chain_ndims,
      shard_axis_names=None,
      name=None,
  ):
      """Helper to `kernel` which computes the log acceptance-correction.
  
      A sufficient but not necessary condition for the existence of a stationary
      distribution, `p(x)`, is "detailed balance", i.e.:
  
      ```none
      p(x'|x) p(x) = p(x|x') p(x')
      ```
  
      In the Metropolis-Hastings algorithm, a state is proposed according to
      `g(x'|x)` and accepted according to `a(x'|x)`, hence
      `p(x'|x) = g(x'|x) a(x'|x)`.
  
      Inserting this into the detailed balance equation implies:
  
      ```none
          g(x'|x) a(x'|x) p(x) = g(x|x') a(x|x') p(x')
      ==> a(x'|x) / a(x|x') = p(x') / p(x) [g(x|x') / g(x'|x)]    (*)
      ```
  
      One definition of `a(x'|x)` which satisfies (*) is:
  
      ```none
      a(x'|x) = min(1, p(x') / p(x) [g(x|x') / g(x'|x)])
      ```
  
      (To see that this satisfies (*), notice that under this definition only at
      most one `a(x'|x)` and `a(x|x') can be other than one.)
  
      We call the bracketed term the "acceptance correction".
  
      In the case of UncalibratedHMC, the log acceptance-correction is not the log
      proposal-ratio. UncalibratedHMC augments the state-space with momentum, z.
      Assuming a standard Gaussian distribution for momentums, the chain eventually
      converges to:
  
      ```none
      p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
      ```
  
      Relating this back to Metropolis-Hastings parlance, for HMC we have:
  
      ```none
      p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
      g([x, z] | [x', z']) = g([x', z'] | [x, z])
      ```
  
      In other words, the MH bracketed term is `1`. However, because we desire to
      use a general MH framework, we can place the momentum probability ratio inside
      the metropolis-correction factor thus getting an acceptance probability:
  
      ```none
                           target_prob(x')
      accept_prob(x'|x) = -----------------  [exp(-0.5 z**2) / exp(-0.5 z'**2)]
                           target_prob(x)
      ```
  
      (Note: we actually need to handle the kinetic energy change at each leapfrog
      step, but this is the idea.)
  
      Args:
        current_momentums: `Tensor` representing the value(s) of the current
          momentum(s) of the state (parts).
        proposed_momentums: `Tensor` representing the value(s) of the proposed
          momentum(s) of the state (parts).
        independent_chain_ndims: Scalar `int` `Tensor` representing the number of
          leftmost `Tensor` dimensions which index independent chains.
        shard_axis_names: A structure of string names indicating how
          members of the state are sharded.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., 'compute_log_acceptance_correction').
  
      Returns:
        log_acceptance_correction: `Tensor` representing the `log`
          acceptance-correction.  (See docstring for mathematical definition.)
      """
      with tf.name_scope(name or "compute_log_acceptance_correction"):
  
          def compute_sum_sq(v, shard_axes):
              sum_sq = tf.reduce_sum(
                  v**2.0,
                  axis=ps.range(independent_chain_ndims, ps.rank(v)),
              )
              if shard_axes is not None:
                  sum_sq = distribute_lib.psum(sum_sq, shard_axes)
              return sum_sq
  
          shard_axis_names = shard_axis_names or (
              [None] * len(current_momentums)
          )
          current_kinetic = tf.add_n(
              [
                  compute_sum_sq(v, axes)
                  for v, axes in zip(
                      current_momentums, shard_axis_names
                  )
              ]
          )
          proposed_kinetic = tf.add_n(
              [
                  compute_sum_sq(v, axes)
                  for v, axes in zip(
                      proposed_momentums, shard_axis_names
                  )
              ]
          )
          return 0.5 * mcmc_util.safe_sum(
              [current_kinetic, -proposed_kinetic]
          )
  
imports: import collections
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 2 for calling a global function with example usage: _compute_log_acceptance_correction(current_momentums, proposed_momentums, independent_chain_ndims, shard_axis_names, name) and returns: unknown, sum_sq