code: def testAdaptingPerReplicaStepSize(self, use_untempered_lp, use_xla):
      num_chains, num_events = 4, 2
      num_results = 250
      target = mvn_diag.MultivariateNormalDiag(loc=[0.0] * num_events)
  
      # The inverse_temperatures are decaying rapidly enough to force
      # significantly different adapted step sizes.
      # It's important not to make inverse_temperatures[-1] too small, or else
      # the last few temperatures could be essentially the same when using an
      # untempered_log_prob_fn.
      inverse_temperatures = 0.25 ** np.arange(3, dtype=np.float32)
      num_replica = len(inverse_temperatures)
  
      # step_size is...
      #   Too large == 3!
      #   A shape that will broadcast across the chains, and (after adaptation)
      #   give a different value for each replica.
      initial_step_size = 3 * np.ones(
          (num_replica, 1, 1), dtype=np.float32
      )
      target_accept_prob = 0.9
  
      def _get_states_and_trace():
          # Whether using the untempered_log_prob_fn setup or not, the target is
          # target.log_prob.
          # However, the untempered_log_prob_fn is very dispersed...it is as though
          # it has a temperature of 100.
          if use_untempered_lp:
              target_log_prob_fn = None
              untempered_log_prob_fn = (
                  lambda x: target.log_prob(x) / 100.0
              )
              tempered_log_prob_fn = (
                  lambda x: target.log_prob(x) * 99 / 100.0
              )
          else:
              target_log_prob_fn = target.log_prob
              untempered_log_prob_fn = None
              tempered_log_prob_fn = None
  
          def make_kernel_fn(target_log_prob_fn):
              kernel = hmc.HamiltonianMonteCarlo(
                  target_log_prob_fn=target_log_prob_fn,
                  step_size=initial_step_size,
                  num_leapfrog_steps=3,
              )
              return ssa.SimpleStepSizeAdaptation(
                  inner_kernel=kernel,
                  target_accept_prob=target_accept_prob,
                  adaptation_rate=0.01,
                  num_adaptation_steps=num_results,
              )
  
          adaptive_remc = replica_exchange_mc.ReplicaExchangeMC(
              target_log_prob_fn=target_log_prob_fn,
              untempered_log_prob_fn=untempered_log_prob_fn,
              tempered_log_prob_fn=tempered_log_prob_fn,
              inverse_temperatures=inverse_temperatures,
              make_kernel_fn=make_kernel_fn,
              state_includes_replicas=True,
          )
  
          def trace_fn(_, results):
              step_adapt_results = results.post_swap_replica_results
              hmc_results = step_adapt_results.inner_results
              return {
                  "step_size": step_adapt_results.new_step_size,
                  "prob_accept": tf.math.exp(
                      tf.math.minimum(hmc_results.log_accept_ratio, 0.0)
                  ),
              }
  
          return sample.sample_chain(
              num_results=num_results,
              current_state=tf.zeros(
                  (num_replica, num_chains, num_events)
              ),
              kernel=adaptive_remc,
              trace_fn=trace_fn,
              num_burnin_steps=0,
              seed=_set_seed(),
          )
  
      if use_xla:
          states, trace = tf.function(
              _get_states_and_trace, jit_compile=True
          )()
      else:
          states, trace = _get_states_and_trace()
  
      states_, final_step_size_, mean_accept_, ess_ = self.evaluate(
          [
              states,
              trace["step_size"][-1],
              tf.reduce_mean(
                  trace["prob_accept"][num_results // 2 :], axis=0
              ),
              effective_sample_size(states, cross_chain_dims=-2),
          ]
      )
  
      self.assertAllEqual((num_replica, 1, 1), final_step_size_.shape)
  
      self.assertAllClose(
          target_accept_prob * np.ones_like(mean_accept_),
          mean_accept_,
          atol=0.2,
      )
  
      # Step size should be increasing (with temperature).
      np.testing.assert_array_less(
          0.0, np.diff(final_step_size_.ravel())
      )
  
      # Step size for the Temperature = 1 replica should have decreased
      # significantly from the large initial value.
      self.assertEqual(initial_step_size[0, 0, 0], 3)
      self.assertLess(final_step_size_[0, 0, 0], 2)
  
      # The mean shouldn't be ridiculous.
      for r in range(num_replica):
          x = states_[:, r]
          mean = np.mean(x)
          stddev = np.std(x)
          n = np.min(ess_[r])
          self.assertAllClose(
              np.zeros_like(mean),
              mean,
              atol=4 * stddev / np.sqrt(n),
              msg="Failed at replica {}".format(r),
          )
  
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 4 for calling the REMCTest obj's underlying member method code with example usage: obj.testAdaptingPerReplicaStepSize(use_untempered_lp, use_xla) and returns: sample, ssa, dict