code: def default_swap_proposal_fn(prob_swap, name=None):
      """Make the default swap proposal func, with `P[swap]`, for replica swap MC.
  
      With probability `prob_swap`, propose combinations of replicas to swap
      When exchanging, create combinations of adjacent replicas in
      [Replica Exchange Monte Carlo](
      https://en.wikipedia.org/wiki/Parallel_tempering).  See also review paper [1].
  
      ```
      swap_fn = default_swap_proposal_fn(prob_swap=0.5)
  
      swap_fn(num_replica=3)
      ==> [1, 0, 2]  # 1 swap, 0 <--> 1
  
      swap_fn(num_replica=3)
      ==> [0, 1, 2]  # 0 swaps
  
      swap_fn(num_replica=3, batch_shape=[2])
      ==> [[0, 1],
           [2, 0],
           [1, 2]]
      ```
  
      Args:
        prob_swap: Scalar `Tensor` in `[0, 1]` giving probability that any swaps
          will be generated.
        name: Python `str` name given to ops created by this function.
          Default value: `'adjacent_swaps'`.
  
      Returns:
        default_swap_proposal_fn_: Python callable which take a number of
          replicas (a Python integer), and integer `Tensor` `batch_shape`, an
          unused `step_count`, a `seed`, and returns `swaps`, a shape
          `[num_replica] + batch_shape` `Tensor`, where axis 0 indexes
          "one-time swaps", i.e., such that (if `rank(swaps) == 1`,
          `range(num_replicas) == tf.gather(swaps, swaps)`.
  
      #### References
  
      [1]: David J. Earl, Michael W. Deem
           Parallel Tempering: Theory, Applications, and New Perspectives
           https://arxiv.org/abs/physics/0508111
      """
  
      def adjacent_swaps(
          num_replica, batch_shape=(), step_count=None, seed=None
      ):
          """Make random shuffle using only one time swaps."""
          del step_count  # Unused for this function.
          with tf.name_scope(name or "adjacent_swaps"):
              parity_seed, proposal_seed = samplers.split_seed(seed)
              # u selects parity.  E.g.,
              #  u==False ==> [1, 0, 3, 2, 4] even parity swaps
              #  u==True ==>  [0, 2, 1, 4, 3] odd parity swaps
              # If there are only 2 replicas, then the "True" swaps are null
              # swaps...which would contradict the user provided `prob_swap`.
              # So special case num_replica==2, forcing u==False in this case.
              u_shape = ps.concat(
                  (
                      ps.ones(1, dtype=tf.int32),
                      ps.cast(batch_shape, tf.int32),
                  ),
                  axis=0,
              )
              u = samplers.uniform(u_shape, seed=parity_seed) < 0.5
              u = tf.where(num_replica > 2, u, False)
  
              x = bu.left_justified_expand_dims_to(
                  ps.range(num_replica, dtype=tf.int64),
                  rank=ps.size(u_shape),
              )
              y = tf.where(
                  tf.equal(x % 2, tf.cast(u, dtype=tf.int64)),
                  x + 1,
                  x - 1,
              )
              y = tf.clip_by_value(y, 0, num_replica - 1)
              # TODO(b/142689785): Consider using tf.cond and returning an empty list
              # then in REMC consider using a tf.cond for short-circuiting.
              return tf.where(
                  samplers.uniform(batch_shape, seed=proposal_seed)
                  < prob_swap,
                  y,
                  x,
              )
  
      return adjacent_swaps
  
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 1 for calling a global function with example usage: default_swap_proposal_fn(prob_swap, name) and returns: adjacent_swaps, tf