class_docstr: Runs one step of the Replica Exchange Monte Carlo.
  
  [Replica Exchange Monte Carlo](
  https://en.wikipedia.org/wiki/Parallel_tempering) is a Markov chain
  Monte Carlo (MCMC) algorithm that is also known as Parallel Tempering. This
  algorithm takes multiple samples (from tempered distributions) in parallel,
  then swaps these samples according to the Metropolis-Hastings criterion.
  See also the review paper [1].
  
  The `K` replicas are parameterized in terms of `inverse_temperature`'s,
  `(beta[0], beta[1], ..., beta[K-1])`. If the user provides
  `target_log_prob_fn`, then the `kth` replica samples from density `p_k(x)`,
  with `log(p_k(x)) = beta_k * target_log_prob(x)`.
  In this case, geometrically decaying `beta` often works well.  That is, with
  `R < 1`, we recommend trying `beta[k] = R^k` so that
  `1.0 = beta[0] > beta[1] > ... > 0`.  See [2].
  
  The user can also provide two functions, `tempered_log_prob_fn` and
  `untempered_log_prob_fn`. In this case, the `kth` replica samples from
  density `p_k(x)` with
  `log(p_k(x)) = beta_k * tempered_log_prob_fn(x) + untempered_log_prob_fn(x)`.
  In this case, `beta` may be zero, and one often sets `beta[-1]` to zero.
  This means the last replica samples using `untempered_log_prob_fn`.
  In the Bayesian setup, `untempered_log_prob_fn` will often be the log prior,
  and `tempered_log_prob_fn` the likelihood.
  
  In all cases,
  
  * `beta[0] == 1` ==> First replica samples from the target density.
  * `beta[k] < 1`, for `k = 1, ..., K-1` ==> Other replicas sample from
    "tempered" versions of target (peak is less high, valley less low).  These
    distributions should allow easier exploration of separated modes.
  
  By default, samples from adjacent replicas `i`, `i + 1` are used as proposals
  for each other in a Metropolis step.  This allows the lower `beta` samples,
  which explore less dense areas of `p`, to eventually swap state with the
  `beta == 1` chain, allowing it to explore these new regions.
  
  Samples from replica 0 are returned, and the others are discarded, unless
  `state_includes_replicas`.
  
  #### Examples
  
  ##### Sampling from the Standard Normal Distribution.
  
  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions
  
  dtype = np.float32
  
  target = tfd.Normal(loc=dtype(0), scale=dtype(1))
  
  # Geometric decay is a good rule of thumb.
  inverse_temperatures = 0.5**tf.range(4, dtype=dtype)
  
  # If everything was Normal, step_size should be ~ sqrt(temperature).
  step_size = 1.5 / tf.sqrt(inverse_temperatures)
  
  def make_kernel_fn(target_log_prob_fn):
    return tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        step_size=step_size, num_leapfrog_steps=3)
  
  remc = tfp.mcmc.ReplicaExchangeMC(
      target_log_prob_fn=target.log_prob,
      inverse_temperatures=inverse_temperatures,
      make_kernel_fn=make_kernel_fn)
  
  def trace_swaps(unused_state, results):
    return (results.is_swap_proposed_adjacent,
            results.is_swap_accepted_adjacent)
  
  samples, (is_swap_proposed_adjacent, is_swap_accepted_adjacent) = (
      tfp.mcmc.sample_chain(
          num_results=1000,
          current_state=1.0,
          kernel=remc,
          num_burnin_steps=500,
          trace_fn=trace_swaps)
  )
  
  # conditional_swap_prob[k] = P[ExchangeAccepted | ExchangeProposed],
  # for the swap between replicas k and k+1.
  conditional_swap_prob = (
      tf.reduce_sum(tf.cast(is_swap_accepted_adjacent, tf.float32), axis=0)
      /
      tf.reduce_sum(tf.cast(is_swap_proposed_adjacent, tf.float32), axis=0))
  ```
  
  ##### Sampling from a 2-D Mixture Normal Distribution.
  
  ```python
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  import matplotlib.pyplot as plt
  tfd = tfp.distributions
  
  dtype = np.float32
  
  target = tfd.MixtureSameFamily(
      mixture_distribution=tfd.Categorical(probs=[0.5, 0.5]),
      components_distribution=tfd.MultivariateNormalDiag(
          loc=[[-1., -1], [1., 1.]],
          scale_diag=0.1*tf.ones([2, 2])))
  
  inverse_temperatures = 0.2**tf.range(4, dtype=dtype)
  
  # step_size must broadcast with all batch and event dimensions of target.
  # Here, this means it must broadcast with:
  #  [len(inverse_temperatures)] + target.event_shape
  step_size = 0.075 / tf.reshape(tf.sqrt(inverse_temperatures), shape=(4, 1))
  
  def make_kernel_fn(target_log_prob_fn):
    return tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        step_size=step_size, num_leapfrog_steps=3)
  
  remc = tfp.mcmc.ReplicaExchangeMC(
      target_log_prob_fn=target.log_prob,
      inverse_temperatures=inverse_temperatures,
      make_kernel_fn=make_kernel_fn)
  
  samples = tfp.mcmc.sample_chain(
      num_results=1000,
      # Start near the [1, 1] mode. Standard HMC would get stuck there.
      current_state=tf.ones(2, dtype=dtype),
      kernel=remc,
      trace_fn=None,
      num_burnin_steps=500)
  
  plt.figure(figsize=(8, 8))
  plt.xlim(-2, 2)
  plt.ylim(-2, 2)
  plt.plot(samples[:, 0], samples[:, 1], '.')
  plt.show()
  ```
  
  #### References
  
  [1]: David J. Earl, Michael W. Deem
       Parallel Tempering: Theory, Applications, and New Perspectives
       https://arxiv.org/abs/physics/0508111
  [2]: David A. Kofke
       On the acceptance probability of replica-exchange Monte Carlo trials.
       J. of Chem. Phys. Vol. 117 No. 5.
class_docstr_tok: ["Runs","one","step","of","the","Replica","Exchange","Monte","Carlo",".","[","Replica","Exchange","Monte","Carlo","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Parallel_tempering",")","is","a","Markov","chain","Monte","Carlo","(","MCMC",")","algorithm","that","is","also","known","as","Parallel","Tempering",".","This","algorithm","takes","multiple","samples","(","from","tempered","distributions",")","in","parallel",",","then","swaps","these","samples","according","to","the","Metropolis-Hastings","criterion",".","See","also","the","review","paper","[","1","]",".","The","`","K","`","replicas","are","parameterized","in","terms","of","`","inverse_temperature","`","'s",",","`","(","beta","[","0","]",",","beta","[","1","]",",","...",",","beta","[","K-1","]",")","`",".","If","the","user","provides","`","target_log_prob_fn","`",",","then","the","`","kth","`","replica","samples","from","density","`","p_k","(","x",")","`",",","with","`","log","(","p_k","(","x",")",")","=","beta_k","*","target_log_prob","(","x",")","`",".","In","this","case",",","geometrically","decaying","`","beta","`","often","works","well",".","That","is",",","with","`","R","<","1","`",",","we","recommend","trying","`","beta","[","k","]","=","R^k","`","so","that","`","1.0","=","beta","[","0","]",">","beta","[","1","]",">","...",">","0","`",".","See","[","2","]",".","The","user","can","also","provide","two","functions",",","`","tempered_log_prob_fn","`","and","`","untempered_log_prob_fn","`",".","In","this","case",",","the","`","kth","`","replica","samples","from","density","`","p_k","(","x",")","`","with","`","log","(","p_k","(","x",")",")","=","beta_k","*","tempered_log_prob_fn","(","x",")","+","untempered_log_prob_fn","(","x",")","`",".","In","this","case",",","`","beta","`","may","be","zero",",","and","one","often","sets","`","beta","[","-1","]","`","to","zero",".","This","means","the","last","replica","samples","using","`","untempered_log_prob_fn","`",".","In","the","Bayesian","setup",",","`","untempered_log_prob_fn","`","will","often","be","the","log","prior",",","and","`","tempered_log_prob_fn","`","the","likelihood",".","In","all","cases",",","*","`","beta","[","0","]","==","1","`","==",">","First","replica","samples","from","the","target","density",".","*","`","beta","[","k","]","<","1","`",",","for","`","k","=","1",",","...",",","K-1","`","==",">","Other","replicas","sample","from","``","tempered","''","versions","of","target","(","peak","is","less","high",",","valley","less","low",")",".","These","distributions","should","allow","easier","exploration","of","separated","modes",".","By","default",",","samples","from","adjacent","replicas","`","i","`",",","`","i","+","1","`","are","used","as","proposals","for","each","other","in","a","Metropolis","step",".","This","allows","the","lower","`","beta","`","samples",",","which","explore","less","dense","areas","of","`","p","`",",","to","eventually","swap","state","with","the","`","beta","==","1","`","chain",",","allowing","it","to","explore","these","new","regions",".","Samples","from","replica","0","are","returned",",","and","the","others","are","discarded",",","unless","`","state_includes_replicas","`",".","#","#","#","#","Examples","#","#","#","#","#","Sampling","from","the","Standard","Normal","Distribution",".","``","`","python","import","numpy","as","np","import","tensorflow","as","tf","import","tensorflow_probability","as","tfp","tfd","=","tfp.distributions","dtype","=","np.float32","target","=","tfd.Normal","(","loc=dtype","(","0",")",",","scale=dtype","(","1",")",")","#","Geometric","decay","is","a","good","rule","of","thumb",".","inverse_temperatures","=","0.5","*","*","tf.range","(","4",",","dtype=dtype",")","#","If","everything","was","Normal",",","step_size","should","be","~","sqrt","(","temperature",")",".","step_size","=","1.5","\/","tf.sqrt","(","inverse_temperatures",")","def","make_kernel_fn","(","target_log_prob_fn",")",":","return","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=target_log_prob_fn",",","step_size=step_size",",","num_leapfrog_steps=3",")","remc","=","tfp.mcmc.ReplicaExchangeMC","(","target_log_prob_fn=target.log_prob",",","inverse_temperatures=inverse_temperatures",",","make_kernel_fn=make_kernel_fn",")","def","trace_swaps","(","unused_state",",","results",")",":","return","(","results.is_swap_proposed_adjacent",",","results.is_swap_accepted_adjacent",")","samples",",","(","is_swap_proposed_adjacent",",","is_swap_accepted_adjacent",")","=","(","tfp.mcmc.sample_chain","(","num_results=1000",",","current_state=1.0",",","kernel=remc",",","num_burnin_steps=500",",","trace_fn=trace_swaps",")",")","#","conditional_swap_prob","[","k","]","=","P","[","ExchangeAccepted","|","ExchangeProposed","]",",","#","for","the","swap","between","replicas","k","and","k+1",".","conditional_swap_prob","=","(","tf.reduce_sum","(","tf.cast","(","is_swap_accepted_adjacent",",","tf.float32",")",",","axis=0",")","\/","tf.reduce_sum","(","tf.cast","(","is_swap_proposed_adjacent",",","tf.float32",")",",","axis=0",")",")","``","`","#","#","#","#","#","Sampling","from","a","2-D","Mixture","Normal","Distribution",".","``","`","python","import","numpy","as","np","import","tensorflow","as","tf","import","tensorflow_probability","as","tfp","import","matplotlib.pyplot","as","plt","tfd","=","tfp.distributions","dtype","=","np.float32","target","=","tfd.MixtureSameFamily","(","mixture_distribution=tfd.Categorical","(","probs=","[","0.5",",","0.5","]",")",",","components_distribution=tfd.MultivariateNormalDiag","(","loc=","[","[","-1.",",","-1","]",",","[","1.",",","1",".","]","]",",","scale_diag=0.1","*","tf.ones","(","[","2",",","2","]",")",")",")","inverse_temperatures","=","0.2","*","*","tf.range","(","4",",","dtype=dtype",")","#","step_size","must","broadcast","with","all","batch","and","event","dimensions","of","target",".","#","Here",",","this","means","it","must","broadcast","with",":","#","[","len","(","inverse_temperatures",")","]","+","target.event_shape","step_size","=","0.075","\/","tf.reshape","(","tf.sqrt","(","inverse_temperatures",")",",","shape=","(","4",",","1",")",")","def","make_kernel_fn","(","target_log_prob_fn",")",":","return","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=target_log_prob_fn",",","step_size=step_size",",","num_leapfrog_steps=3",")","remc","=","tfp.mcmc.ReplicaExchangeMC","(","target_log_prob_fn=target.log_prob",",","inverse_temperatures=inverse_temperatures",",","make_kernel_fn=make_kernel_fn",")","samples","=","tfp.mcmc.sample_chain","(","num_results=1000",",","#","Start","near","the","[","1",",","1","]","mode",".","Standard","HMC","would","get","stuck","there",".","current_state=tf.ones","(","2",",","dtype=dtype",")",",","kernel=remc",",","trace_fn=None",",","num_burnin_steps=500",")","plt.figure","(","figsize=","(","8",",","8",")",")","plt.xlim","(","-2",",","2",")","plt.ylim","(","-2",",","2",")","plt.plot","(","samples","[",":",",","0","]",",","samples","[",":",",","1","]",",","'",".","'",")","plt.show","(",")","``","`","#","#","#","#","References","[","1","]",":","David","J.","Earl",",","Michael","W.","Deem","Parallel","Tempering",":","Theory",",","Applications",",","and","New","Perspectives","https",":","\/\/arxiv.org\/abs\/physics\/0508111","[","2","]",":","David","A.","Kofke","On","the","acceptance","probability","of","replica-exchange","Monte","Carlo","trials",".","J.","of","Chem",".","Phys",".","Vol",".","117","No",".","5","."]
docstr: Takes one step of the TransitionKernel.
  
  Args:
    current_state: `Tensor` or Python `list` of `Tensor`s representing the
      current state(s) of the Markov chain(s).
    previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
      `list` of `Tensor`s representing internal calculations made within the
      previous call to this function (or as returned by `bootstrap_results`).
    seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
  Returns:
    next_state: `Tensor` or Python `list` of `Tensor`s representing the
      next state(s) of the Markov chain(s).
    kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
      `Tensor`s representing internal calculations made within this function.
      This inculdes replica states.
docstr_tok: ["Takes","one","step","of","the","TransitionKernel",".","Args",":","current_state",":","`","Tensor","`","or","Python","`","list","`","of","`","Tensor","`","s","representing","the","current","state","(","s",")","of","the","Markov","chain","(","s",")",".","previous_kernel_results",":","A","(","possibly","nested",")","`","tuple","`",",","`","namedtuple","`","or","`","list","`","of","`","Tensor","`","s","representing","internal","calculations","made","within","the","previous","call","to","this","function","(","or","as","returned","by","`","bootstrap_results","`",")",".","seed",":","PRNG","seed",";","see","`","tfp.random.sanitize_seed","`","for","details",".","Returns",":","next_state",":","`","Tensor","`","or","Python","`","list","`","of","`","Tensor","`","s","representing","the","next","state","(","s",")","of","the","Markov","chain","(","s",")",".","kernel_results",":","A","(","possibly","nested",")","`","tuple","`",",","`","namedtuple","`","or","`","list","`","of","`","Tensor","`","s","representing","internal","calculations","made","within","this","function",".","This","inculdes","replica","states","."]
code: def one_step(self, current_state, previous_kernel_results, seed=None):
      """Takes one step of the TransitionKernel.
  
      Args:
        current_state: `Tensor` or Python `list` of `Tensor`s representing the
          current state(s) of the Markov chain(s).
        previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
          `list` of `Tensor`s representing internal calculations made within the
          previous call to this function (or as returned by `bootstrap_results`).
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        next_state: `Tensor` or Python `list` of `Tensor`s representing the
          next state(s) of the Markov chain(s).
        kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
          `Tensor`s representing internal calculations made within this function.
          This inculdes replica states.
      """
  
      # The code below propagates one step states of shape
      #  [n_replica] + batch_shape + event_shape.
      #
      # The step is done in three parts:
      #  1) Call one_step to transition states via a tempered version of
      #     self.target_log_prob_fn (see _replica_target_log_prob).
      #  2) Permute values in states
      #  3) Update state-dependent values, such as log_probs.
      #
      # We chose to swap states, rather than temperatures, because...
      # (i)  If swapping temperatures, you *still* have to swap log_probs to
      #      determine acceptance, as well as states (for kernel results).
      #      So it's just as difficult to swap temperatures.
      # (ii) If swapping temperatures, you have to take care to swap any user-
      #      supplied temperature related things (like step size).
      #      A-priori, we don't know what else will need to be swapped!
      # (iii)In both cases, the kernel results need to be updated in a non-trivial
      #      manner....so we either special-case, or use bootstrap.
  
      with tf.name_scope(
          mcmc_util.make_name(self.name, "remc", "one_step")
      ):
          # Force a read in case the `inverse_temperatures` is a `tf.Variable`.
          inverse_temperatures = tf.convert_to_tensor(
              previous_kernel_results.inverse_temperatures,
              name="inverse_temperatures",
          )
  
          target_log_prob_for_inner_kernel = (
              _make_replica_target_log_prob_fn(
                  target_log_prob_fn=self.target_log_prob_fn,
                  inverse_temperatures=inverse_temperatures,
                  untempered_log_prob_fn=self.untempered_log_prob_fn,
                  tempered_log_prob_fn=self.tempered_log_prob_fn,
              )
          )
          # TODO(b/159636942): Clean up the helpful error msg after 2020-11-10.
          try:
              inner_kernel = (
                  self.make_kernel_fn(  # pylint: disable=not-callable
                      target_log_prob_for_inner_kernel
                  )
              )
          except TypeError as e:
              if "argument" not in str(e):
                  raise
              raise TypeError(
                  "`ReplicaExchangeMC`s `make_kernel_fn` no longer receives a `seed` "
                  "argument. `TransitionKernel` instances now receive seeds via "
                  "`one_step`."
              )
  
          seed = samplers.sanitize_seed(seed)  # Retain for diagnostics.
          inner_seed, swap_seed, logu_seed = samplers.split_seed(
              seed, n=3
          )
          # Step the inner TransitionKernel.
          [
              pre_swap_replica_states,
              pre_swap_replica_results,
          ] = inner_kernel.one_step(
              previous_kernel_results.post_swap_replica_states,
              previous_kernel_results.post_swap_replica_results,
              seed=inner_seed,
          )
  
          pre_swap_replica_target_log_prob = _get_field(
              # These are tempered log probs (have been divided by temperature).
              pre_swap_replica_results,
              "target_log_prob",
          )
  
          dtype = pre_swap_replica_target_log_prob.dtype
          replica_and_batch_shape = ps.shape(
              pre_swap_replica_target_log_prob
          )
          batch_shape = replica_and_batch_shape[1:]
          replica_and_batch_rank = ps.rank(
              pre_swap_replica_target_log_prob
          )
          num_replica = ps.size0(inverse_temperatures)
  
          inverse_temperatures = bu.left_justified_broadcast_to(
              inverse_temperatures, replica_and_batch_shape
          )
  
          # Now that each replica has done one_step, it is time to consider swaps.
  
          # swap.shape = [n_replica], and is a "once only" permutation, meaning it
          # is achievable by a sequence of pairwise permutations, where each element
          # is moved at most once.
          # E.g. if swaps = [1, 0, 2], we will consider swapping temperatures 0 and
          # 1, keeping 2 fixed.  This exact same swap is considered for *every*
          # batch member.  Of course some batch members may accept and some reject.
          try:
              swaps = tf.cast(
                  self.swap_proposal_fn(  # pylint: disable=not-callable
                      num_replica,
                      batch_shape=batch_shape,
                      seed=swap_seed,
                      step_count=previous_kernel_results.step_count,
                  ),
                  dtype=tf.int32,
              )
          except TypeError as e:
              if "step_count" not in str(e):
                  raise
              warnings.warn(
                  "The `swap_proposal_fn` given to ReplicaExchangeMC did not accept "
                  "the `step_count` argument. Falling back to omitting the "
                  "argument. This fallback will be removed after 24-Oct-2020."
              )
              swaps = tf.cast(
                  self.swap_proposal_fn(  # pylint: disable=not-callable
                      num_replica,
                      batch_shape=batch_shape,
                      seed=swap_seed,
                  ),
                  dtype=tf.int32,
              )
  
          null_swaps = bu.left_justified_expand_dims_like(
              tf.range(num_replica, dtype=swaps.dtype), swaps
          )
          swaps = _maybe_embed_swaps_validation(
              swaps, null_swaps, self.validate_args
          )
  
          # Un-temper the log probs for use in the swap acceptance ratio.
          if self.tempered_log_prob_fn is None:
              # Efficient way of re-evaluating target_log_prob_fn on the
              # pre_swap_replica_states.
              untempered_negative_energy_ignoring_ulp = (
                  # Since untempered_log_prob_fn is None, we may assume
                  # inverse_temperatures > 0 (else the target is improper).
                  pre_swap_replica_target_log_prob
                  / inverse_temperatures
              )
          else:
              # The untempered_log_prob_fn does not factor into the acceptance ratio.
              # Proof: Suppose the tempered target is
              #   p_k(x) = f(x)^{beta_k} g(x),
              # So f(x) is tempered, and g(x) is not.  Then, the acceptance ratio for
              # a 1 <--> 2 swap is...
              #   (p_1(x_2) p_2(x_1)) / (p_1(x_1) p_2(x_2))
              # which depends only on f(x), since terms involving g(x) cancel.
              untempered_negative_energy_ignoring_ulp = (
                  self.tempered_log_prob_fn(*pre_swap_replica_states)
              )
  
          # Since `swaps` is its own inverse permutation we automatically know the
          # swap counterpart: range(num_replica). We use this idea to compute the
          # acceptance in a vectorized manner at the cost of wasting roughly half
          # our computation. Although we could use `unique` to solve this problem,
          # we expect the cost of `unique` to be higher than the dozens of wasted
          # arithmetic calculations. Worse, it'd mean we need dynamic sized Tensors
          # (eg, using `tf.where(bool)`) and so we wouldn't be able to XLA compile.
  
          # Note: diffs would normally be "proposed - current" however energy is
          # flipped since `energy == -log_prob`.
          # Note: The untempered_log_prob_fn (if provided) is not included in
          # untempered_pre_swap_replica_target_log_prob, and hence does not factor
          # into energy_diff. Why? Because, it cancels out in the acceptance ratio.
          energy_diff = (
              untempered_negative_energy_ignoring_ulp
              - mcmc_util.index_remapping_gather(
                  untempered_negative_energy_ignoring_ulp,
                  swaps,
                  name="gather_swap_tlp",
              )
          )
          swapped_inverse_temperatures = (
              mcmc_util.index_remapping_gather(
                  inverse_temperatures, swaps, name="gather_swap_temps"
              )
          )
          inverse_temp_diff = (
              swapped_inverse_temperatures - inverse_temperatures
          )
  
          # If i and j are swapping, log_accept_ratio[] i and j are equal.
          log_accept_ratio = (
              energy_diff
              * bu.left_justified_expand_dims_to(
                  inverse_temp_diff, replica_and_batch_rank
              )
          )
  
          log_accept_ratio = tf.where(
              tf.math.is_finite(log_accept_ratio),
              log_accept_ratio,
              tf.constant(-np.inf, dtype=dtype),
          )
  
          # Produce log[Uniform] draws that are identical at swapped indices.
          log_uniform = tf.math.log(
              samplers.uniform(
                  shape=replica_and_batch_shape,
                  dtype=dtype,
                  seed=logu_seed,
              )
          )
          anchor_swaps = tf.minimum(swaps, null_swaps)
          log_uniform = mcmc_util.index_remapping_gather(
              log_uniform, anchor_swaps
          )
  
          is_swap_accepted_mask = tf.less(
              log_uniform,
              log_accept_ratio,
              name="is_swap_accepted_mask",
          )
  
          def _swap_tensor(x):
              return mcmc_util.choose(
                  is_swap_accepted_mask,
                  mcmc_util.index_remapping_gather(x, swaps),
                  x,
              )
  
          post_swap_replica_states = [
              _swap_tensor(s) for s in pre_swap_replica_states
          ]
  
          expanded_null_swaps = bu.left_justified_broadcast_to(
              null_swaps, replica_and_batch_shape
          )
          is_swap_proposed = _compute_swap_notmatrix(
              # Broadcast both so they have shape [num_replica] + batch_shape.
              # This (i) makes them have same shape as is_swap_accepted, and
              # (ii) keeps shape consistent if someday swaps has a batch shape.
              expanded_null_swaps,
              bu.left_justified_broadcast_to(
                  swaps, replica_and_batch_shape
              ),
          )
  
          # To get is_swap_accepted in ordered position, we use
          # _compute_swap_notmatrix on current and next replica positions.
          post_swap_replica_position = _swap_tensor(expanded_null_swaps)
  
          is_swap_accepted = _compute_swap_notmatrix(
              post_swap_replica_position, expanded_null_swaps
          )
  
          if self._state_includes_replicas:
              post_swap_states = post_swap_replica_states
          else:
              post_swap_states = [
                  s[0] for s in post_swap_replica_states
              ]
  
          post_swap_replica_results = _set_swapped_fields_to_nan(
              _swap_log_prob_and_maybe_grads(
                  pre_swap_replica_results,
                  post_swap_replica_states,
                  inner_kernel,
              )
          )
  
          if mcmc_util.is_list_like(current_state):
              # We *always* canonicalize the states in the kernel results.
              states = post_swap_states
          else:
              states = post_swap_states[0]
  
          post_swap_kernel_results = ReplicaExchangeMCKernelResults(
              post_swap_replica_states=post_swap_replica_states,
              pre_swap_replica_results=pre_swap_replica_results,
              post_swap_replica_results=post_swap_replica_results,
              is_swap_proposed=is_swap_proposed,
              is_swap_accepted=is_swap_accepted,
              is_swap_proposed_adjacent=_sub_diag(is_swap_proposed),
              is_swap_accepted_adjacent=_sub_diag(is_swap_accepted),
              # Store the original pkr.inverse_temperatures in case its a
              # `tf.Variable`.
              inverse_temperatures=previous_kernel_results.inverse_temperatures,
              swaps=swaps,
              step_count=previous_kernel_results.step_count + 1,
              seed=seed,
              potential_energy=-untempered_negative_energy_ignoring_ulp,
          )
  
          return states, post_swap_kernel_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 13 for calling the ReplicaExchangeMC obj's underlying member method code with example usage: obj.one_step(current_state, previous_kernel_results, seed) and returns: states,  post_swap_kernel_results, mcmc_util