code: def _compute_log_acceptance_correction(
      current_state_parts,
      proposed_state_parts,
      current_volatility_parts,
      proposed_volatility_parts,
      current_drift_parts,
      proposed_drift_parts,
      step_size_parts,
      independent_chain_ndims,
      experimental_shard_axis_names=None,
      name=None,
  ):
      r"""Helper to `kernel` which computes the log acceptance-correction.
  
      Computes `log_acceptance_correction` as described in `MetropolisHastings`
      class. The proposal density is normal. More specifically,
  
       ```none
      q(proposed_state | current_state) \sim N(current_state + current_drift,
      step_size * current_volatility**2)
  
      q(current_state | proposed_state) \sim N(proposed_state + proposed_drift,
      step_size * proposed_volatility**2)
      ```
  
      The `log_acceptance_correction` is then
  
      ```none
      log_acceptance_correctio = q(current_state | proposed_state)
      - q(proposed_state | current_state)
      ```
  
      Args:
        current_state_parts: Python `list` of `Tensor`s representing the value(s) of
          the current state of the chain.
        proposed_state_parts:  Python `list` of `Tensor`s representing the value(s)
          of the proposed state of the chain. Must broadcast with the shape of
          `current_state_parts`.
        current_volatility_parts: Python `list` of `Tensor`s representing the value
          of `volatility_fn(*current_volatility_parts)`. Must broadcast with the
          shape of `current_state_parts`.
        proposed_volatility_parts: Python `list` of `Tensor`s representing the value
          of `volatility_fn(*proposed_volatility_parts)`. Must broadcast with the
          shape of `current_state_parts`
        current_drift_parts: Python `list` of `Tensor`s representing value of the
          drift `_get_drift(*current_state_parts, ..)`. Must broadcast with the
          shape of `current_state_parts`.
        proposed_drift_parts: Python `list` of `Tensor`s representing value of the
          drift `_get_drift(*proposed_drift_parts, ..)`. Must broadcast with the
          shape of `current_state_parts`.
        step_size_parts: Python `list` of `Tensor`s representing the step size for
          Euler-Maruyama method. Must broadcast with the shape of
          `current_state_parts`.
        independent_chain_ndims: Scalar `int` `Tensor` representing the number of
          leftmost `Tensor` dimensions which index independent chains.
        experimental_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"):
          proposed_log_density_parts = []
          dual_log_density_parts = []
  
          if experimental_shard_axis_names is None:
              experimental_shard_axis_names = [None] * len(
                  current_state_parts
              )
  
          for [
              current_state,
              proposed_state,
              current_volatility,
              proposed_volatility,
              current_drift,
              proposed_drift,
              step_size,
              shard_axes,
          ] in zip(
              current_state_parts,
              proposed_state_parts,
              current_volatility_parts,
              proposed_volatility_parts,
              current_drift_parts,
              proposed_drift_parts,
              step_size_parts,
              experimental_shard_axis_names,
          ):
              axis = ps.range(
                  independent_chain_ndims, ps.rank(current_state)
              )
  
              state_diff = proposed_state - current_state
  
              current_volatility *= tf.sqrt(step_size)
  
              proposed_energy = (
                  state_diff - current_drift
              ) / current_volatility
  
              proposed_volatility *= tf.sqrt(step_size)
  
              # Compute part of `q(proposed_state | current_state)`
              def reduce_sum(shard_axes, x, axis=None):
                  x = tf.reduce_sum(x, axis)
                  if shard_axes is not None:
                      x = distribute_lib.psum(x, shard_axes)
                  return x
  
              proposed_energy = reduce_sum(
                  shard_axes,
                  mcmc_util.safe_sum(
                      [
                          tf.math.log(current_volatility),
                          0.5 * (proposed_energy**2),
                      ]
                  ),
                  axis=axis,
              )
              proposed_log_density_parts.append(-proposed_energy)
  
              # Compute part of `q(current_state | proposed_state)`
              dual_energy = (
                  state_diff + proposed_drift
              ) / proposed_volatility
              dual_energy = reduce_sum(
                  shard_axes,
                  mcmc_util.safe_sum(
                      [
                          tf.math.log(proposed_volatility),
                          0.5 * (dual_energy**2),
                      ]
                  ),
                  axis=axis,
              )
              dual_log_density_parts.append(-dual_energy)
  
          # Compute `q(proposed_state | current_state)`
          proposed_log_density_reduce = tf.add_n(
              proposed_log_density_parts
          )
          # Compute `q(current_state | proposed_state)`
          dual_log_density_reduce = tf.add_n(dual_log_density_parts)
  
          return mcmc_util.safe_sum(
              [dual_log_density_reduce, -proposed_log_density_reduce]
          )
  
imports: import collections
  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.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 3 for calling a global function with example usage: _compute_log_acceptance_correction(current_state_parts, proposed_state_parts, current_volatility_parts, proposed_volatility_parts, current_drift_parts, proposed_drift_parts, step_size_parts, independent_chain_ndims, experimental_shard_axis_names, name) and returns: mcmc_util, x