docstr: Sampling from two batch diagonal multivariate normal.
code: def test1EventDim2BatchDim3Replica(
      self,
      tfp_transition_kernel,
      inverse_temperatures,
      step_size_fn,
      ess_scaling,
      use_untempered_log_prob_fn,
  ):
      """Sampling from two batch diagonal multivariate normal."""
      step_size = (
          step_size_fn(inverse_temperatures) + np.exp(np.pi) / 100
      ).astype(
          np.float32
      )  # Prevent resonances.
  
      # Small scale and well-separated modes mean we need replica swap to
      # work or else tests fail.
      loc = np.array(
          [
              # Use 3-D normals, ensuring batch and event sizes don't broadcast.
              [-1.0, -0.5, 0.0],  # loc of first batch
              [1.0, 0.5, 0.0],  # loc of second batch
          ],
          dtype=np.float32,
      )
      scale_diag = np.array(
          [[0.5, 0.5, 0.5], [0.8, 0.8, 0.8]], dtype=np.float32
      )
      target = mvn_diag.MultivariateNormalDiag(
          loc=loc, scale_diag=scale_diag
      )
  
      def make_kernel_fn(target_log_prob_fn):
          return tfp_transition_kernel(
              target_log_prob_fn=target_log_prob_fn,
              step_size=step_size,
              num_leapfrog_steps=3,
          )
  
      target_log_prob_fn = tf.function(target.log_prob, autograph=False)
  
      if use_untempered_log_prob_fn:
          # This untempered_log_prob_fn should not change any of the results.
          untempered_log_prob_fn = lambda x: tf.zeros_like(loc[..., 0])
          tempered_log_prob_fn = target_log_prob_fn
          target_log_prob_fn = None
      else:
          untempered_log_prob_fn = None
          tempered_log_prob_fn = None
  
      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,
      )
      remc.one_step = tf.function(remc.one_step, autograph=False)
  
      def trace_fn(state, results):  # pylint: disable=unused-argument
          return [
              unnest.get_innermost(
                  results.post_swap_replica_results, "log_accept_ratio"
              ),
              results.post_swap_replica_states,
          ]
  
      num_results = 2000
      states, (log_accept_ratio, replica_states) = sample.sample_chain(
          num_results=num_results,
          current_state=loc[::-1],  # Batch members far from their mode!
          kernel=remc,
          num_burnin_steps=100,
          trace_fn=trace_fn,
          seed=test_util.test_seed(),
      )
  
      num_replica = inverse_temperatures.shape[0]
  
      self.assertLen(replica_states, 1)  # One state part
      replica_states = replica_states[0]
  
      self.assertAllEqual(
          (num_results, num_replica) + loc.shape, replica_states.shape
      )
      self.assertAllEqual((num_results,) + loc.shape, states.shape)
  
      (
          states_,
          replica_states_,
          replica_mean_,
          replica_cov_,
          accept_probs_,
          ess_,
      ) = self.evaluate(
          [
              states,
              replica_states,
              tf.reduce_mean(replica_states, axis=0),
              sample_stats.covariance(replica_states),
              tf.math.exp(tf.minimum(0.0, log_accept_ratio)),
              effective_sample_size(replica_states),
          ]
      )
  
      logging.vlog(
          2,
          "---- execution:{} Min[ESS]: {}  mean_accept: {}".format(
              "eager" if tf.executing_eagerly() else "graph",
              np.min(ess_),
              np.mean(accept_probs_, axis=0),
          ),
      )
  
      self.assertAllEqual(states_, replica_states_[:, 0])
  
      def _check_stats(replica_idx, batch_idx, ess_scaling):
          err_msg = "Failure in replica {}, batch {}".format(
              replica_idx, batch_idx
          )
          assert inverse_temperatures.ndim in [1, 2]
          if inverse_temperatures.ndim == 1:
              temperature = 1 / inverse_temperatures[replica_idx]
          elif inverse_temperatures.ndim == 2:
              temperature = (
                  1 / inverse_temperatures[replica_idx, batch_idx]
              )
  
          expected_scale = scale_diag[batch_idx, 0] * np.sqrt(
              temperature
          )
  
          ess = np.min(
              ess_[replica_idx, batch_idx]
          )  # Conservative estimate.
          self.assertGreater(
              ess, num_results * ess_scaling, msg="Bad sampling!"
          )
  
          self.assertAllClose(
              replica_mean_[replica_idx, batch_idx],
              loc[batch_idx],
              # 6 standard errors of a mean estimate.
              atol=6 * expected_scale / np.sqrt(ess),
              msg=err_msg,
          )
          self.assertAllClose(
              expected_scale**2 * np.eye(loc.shape[1]),
              replica_cov_[replica_idx, batch_idx],
              # 12 standard errors of a variance estimate.
              atol=12 * np.sqrt(2) * expected_scale**2 / np.sqrt(ess),
              msg=err_msg,
          )
  
      for replica_idx in range(num_replica):
          for batch_idx in range(loc.shape[0]):
              _check_stats(replica_idx, batch_idx, ess_scaling)
  
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 7 for calling the REMCTest obj's underlying member method code with example usage: obj.test1EventDim2BatchDim3Replica(tfp_transition_kernel, inverse_temperatures, step_size_fn, ess_scaling, use_untempered_log_prob_fn) and returns: tfp_transition_kernel, list