docstr: Sampling from a 2-D Mixture Normal Distribution.
code: def test2DMixNormal(
      self,
      tfp_transition_kernel,
      state_includes_replicas,
      swap_proposal_fn_name,
  ):
      """Sampling from a 2-D Mixture Normal Distribution."""
      swap_proposal_fn = {
          "default": replica_exchange_mc.default_swap_proposal_fn(1.0),
          "even_odd": replica_exchange_mc.even_odd_swap_proposal_fn(
              1.0
          ),
      }[swap_proposal_fn_name]
  
      dtype = np.float32
  
      # By symmetry, target has mean [0, 0]
      # Therefore, Var = E[X^2] = E[E[X^2 | c]], where c is the component.
      # Now..., for the first component,
      #   E[X1^2] =  Var[X1] + Mean[X1]^2
      #           =  0.3^2 + 1^2,
      # and similarly for the second. As a result, Var[mixture] = 1.09.
      target = mixture_same_family.MixtureSameFamily(
          mixture_distribution=categorical.Categorical(
              probs=[0.5, 0.5]
          ),
          components_distribution=mvn_diag.MultivariateNormalDiag(
              loc=[[-1.0, -1], [1.0, 1.0]],
              scale_diag=0.3 * tf.ones((2, 2)),
          ),
      )
  
      num_replica = 4
      inverse_temperatures = 10.0 ** tf.linspace(
          start=0.0, stop=-1.0, num=num_replica
      )
      # We need to pad the step_size so it broadcasts against MCMC samples. In
      # this case we have 1 replica dim, 0 batch dims, and 1 event dim hence need
      # to right pad the step_size by one dim (for the event).
      step_size = 0.2 / tf.math.sqrt(
          inverse_temperatures[:, tf.newaxis]
      )
  
      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=5,
          )
  
      remc = replica_exchange_mc.ReplicaExchangeMC(
          target_log_prob_fn=target.log_prob,
          # Verified that test fails if inverse_temperatures = [1.]
          inverse_temperatures=inverse_temperatures,
          swap_proposal_fn=swap_proposal_fn,
          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"
          )
  
      if state_includes_replicas:
          current_state = tf.ones((num_replica, 2), dtype=dtype)
      else:
          current_state = tf.ones(2, dtype=dtype)
  
      num_results = 2000
      states, replica_log_accept_ratio = sample.sample_chain(
          num_results=num_results,
          # Start at one of the modes, in order to make mode jumping necessary
          # if we want to pass test.
          current_state=current_state,
          kernel=remc,
          num_burnin_steps=50,
          trace_fn=trace_fn,
          seed=test_util.test_seed(),
      )
  
      if state_includes_replicas:
          self.assertAllEqual(
              (num_results, num_replica, 2), states.shape
          )
          states = states[:, 0]  # Keep only replica 0
      else:
          self.assertAllEqual((num_results, 2), states.shape)
  
      replica_accept_ratio = tf.reduce_mean(
          tf.math.exp(tf.minimum(0.0, replica_log_accept_ratio)), axis=0
      )
  
      [
          sample_mean_,
          sample_variance_,
          replica_accept_ratio_,
          expected_mean_,
          expected_stddev_,
          expected_variance_,
          ess_,
      ] = self.evaluate(
          [
              tf.reduce_mean(states, axis=0),
              sample_stats.variance(states),
              replica_accept_ratio,
              target.mean(),
              target.stddev(),
              target.variance(),
              effective_sample_size(states),
          ]
      )
  
      logging.vlog(
          2,
          "---- execution:{}  accept_ratio:{}  mean:{}".format(
              "eager" if tf.executing_eagerly() else "graph",
              replica_accept_ratio_,
              sample_mean_,
          ),
      )
  
      mean_atol = 6 * expected_stddev_ / np.sqrt(np.min(ess_))
      var_atol = 6 * expected_variance_ / np.sqrt(np.min(ess_))
      for i in range(mean_atol.shape[0]):
          self.assertAllClose(
              expected_mean_[i],
              sample_mean_[i],
              atol=mean_atol[i],
              msg="position {}".format(i),
          )
          self.assertAllClose(
              expected_variance_[i],
              sample_variance_[i],
              atol=var_atol[i],
              msg=i,
          )
  
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 5 for calling the REMCTest obj's underlying member method code with example usage: obj.test2DMixNormal(tfp_transition_kernel, state_includes_replicas, swap_proposal_fn_name) and returns: tfp_transition_kernel, unnest