code: def testYesBatchThreeReplicas(self):
      # There are three replicas, so we get...
      #   null swaps [0, 1, 2],
      #   even swaps [1, 0, 2],
      #   odd swaps [0, 2, 1].
      # Each type is identically repeated across all batch members.
      fn = replica_exchange_mc.even_odd_swap_proposal_fn(
          swap_frequency=0.5
      )
  
      num_batch = 2
      num_replica = 3
  
      null_swap = [[0] * num_batch, [1] * num_batch, [2] * num_batch]
  
      even_swap = [[1] * num_batch, [0] * num_batch, [2] * num_batch]
  
      odd_swap = [[0] * num_batch, [2] * num_batch, [1] * num_batch]
  
      with self.subTest(
          "Even swap since count=0 is an even multiple of period=2"
      ):
          swaps = fn(
              num_replica=num_replica,
              batch_shape=[num_batch],
              step_count=0,
          )
          self.assertAllEqual((num_replica, num_batch), swaps.shape)
          self.assertAllEqual(even_swap, self.evaluate(swaps))
  
      with self.subTest(
          "No swap since count=1 is not a multiple of period=2"
      ):
          swaps = fn(
              num_replica=num_replica,
              batch_shape=[num_batch],
              step_count=1,
          )
          self.assertAllEqual((num_replica, num_batch), swaps.shape)
          self.assertAllEqual(null_swap, self.evaluate(swaps))
  
      with self.subTest(
          "Odd swap since count=2 is an odd multiple of period=2"
      ):
          swaps = fn(
              num_replica=num_replica,
              batch_shape=[num_batch],
              step_count=2,
          )
          self.assertAllEqual((num_replica, num_batch), swaps.shape)
          self.assertAllEqual(odd_swap, self.evaluate(swaps))
  
      with self.subTest(
          "No swap since count=3 is not a multiple of period=2"
      ):
          swaps = fn(
              num_replica=num_replica,
              batch_shape=[num_batch],
              step_count=3,
          )
          self.assertAllEqual(null_swap, self.evaluate(swaps))
  
      with self.subTest(
          "Even swap since count=4 is an even multiple of period=2"
      ):
          swaps = fn(
              num_replica=num_replica,
              batch_shape=[num_batch],
              step_count=4,
          )
          self.assertAllEqual(even_swap, self.evaluate(swaps))
  
imports: from absl import logging
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_combinations
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import random_walk_metropolis
  from tensorflow_probability.python.mcmc import replica_exchange_mc
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 1 for calling the EvenOddSwapProposedFnTest obj's underlying member method code with example usage: obj.testYesBatchThreeReplicas() without return types