code: def _swap_log_prob_and_maybe_grads(
      pre_swap_replica_results,
      post_swap_replica_states,
      inner_kernel,
  ):
      """Swap 'target_log_prob' and maybe 'grads_target_log_prob'.
  
      After swapping states, the log probs for replicas need to be evaluated at
      the new points, then tempered. The easiest way to do this is to call
      `inner_kernel.bootstrap_results`.
  
      Args:
        pre_swap_replica_results: Replica results, before swapping.
        post_swap_replica_states: The states, after swapping.
        inner_kernel: The tempered inner kernel used by ReplicaExchangeMC.
  
      Returns:
        Copy of pre_swap_replica_results, with some fields set to NaN.
      """
      # This call means we re-evaluate gradients. More efficient surgery could be
      # done, but this is much simpler.
      temp_kr = inner_kernel.bootstrap_results(post_swap_replica_states)
  
      post_swap_replica_results = pre_swap_replica_results
  
      # Fields we expect to find, else raise.
      for field in ["target_log_prob"]:
          post_swap_replica_results = _update_field(
              post_swap_replica_results,
              field,
              _get_field(temp_kr, field),
          )
  
      # Fields we don't always find.
      for field in ["grads_target_log_prob"]:
          try:
              post_swap_replica_results = _update_field(
                  post_swap_replica_results,
                  field,
                  _get_field(temp_kr, field),
              )
          except REMCFieldNotFoundError:
              pass
  
      return post_swap_replica_results
  
imports: import collections
  import warnings
  import numpy
  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 distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 7 for calling a global function with example usage: _swap_log_prob_and_maybe_grads(pre_swap_replica_results, post_swap_replica_states, inner_kernel) and returns: post_swap_replica_results