code: def testProbSwapNumReplicaWithBatch(self, prob_swap, num_replica):
      fn = replica_exchange_mc.default_swap_proposal_fn(prob_swap)
      num_results = 100
      seeds = samplers.split_seed(test_util.test_seed(), n=num_results)
      swaps = tf.stack(
          [
              fn(num_replica, batch_shape=[2], seed=seeds[i])
              for i in range(num_results)
          ],
          axis=0,
      )
  
      self.assertAllEqual((num_results, num_replica, 2), swaps.shape)
      swaps_ = self.evaluate(swaps)
  
      # Batch members should have distinct swaps in most cases.
      frac_same = np.mean(swaps_[..., 0] == swaps_[..., 1])
  
      # If prob_swap == 0, swap is the null_swap always.
      if (
          prob_swap == 0
          or
          # If num_replica == 1, swap = [0] always.
          num_replica == 1
          or
          # In this case, we always swap and it's always [1, 0].
          (num_replica == 2 and prob_swap == 1)
      ):
          self.assertEqual(1.0, frac_same)
      else:
          self.assertLess(frac_same, 0.9)
  
      # Check that each batch member has proper statistics.
      for i in range(swaps_.shape[-1]):
          self.check_swaps_with_no_batch_shape(
              swaps_[..., i], prob_swap
          )
  
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 2 for calling the DefaultSwapProposedFnTest obj's underlying member method code with example usage: obj.testProbSwapNumReplicaWithBatch(prob_swap, num_replica) without return types