code: def testMultipleCorrelatedStatesWithNoBatchDims(self):
      dtype = np.float32
      num_results = 2000
      true_mean = dtype([0, 0])
      true_cov = dtype([[1, 0.5], [0.5, 1]])
      # Use LinearOperatorLowerTriangular to get broadcasting ability.
      linop = tf.linalg.LinearOperatorLowerTriangular(
          tf.linalg.cholesky(true_cov)
      )
  
      # Its ok to decorate this since we only need to stress the TransitionKernel.
      def target_log_prob(x, y):
          # Corresponds to unnormalized MVN.
          # z = matmul(inv(chol(true_cov)), [x, y] - true_mean)
          xy = tf.stack([x, y], axis=-1) - true_mean
          z = linop.solvevec(xy)
          return -0.5 * tf.reduce_sum(z**2.0, axis=-1)
  
      inverse_temperatures = tf.constant([1.0, 0.75, 0.5])
      # 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 0 event dims (per each
      # of 2 state parts) hence no padding is needed.
      # We do however supply a step size for each state part.
      step_sizes = [0.9 / tf.math.sqrt(inverse_temperatures)] * 2
  
      def make_kernel_fn(target_log_prob_fn):
          return hmc.HamiltonianMonteCarlo(
              target_log_prob_fn=target_log_prob_fn,
              step_size=step_sizes,
              num_leapfrog_steps=3,
          )
  
      remc = replica_exchange_mc.ReplicaExchangeMC(
          target_log_prob_fn=target_log_prob,
          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 results.post_swap_replica_results.log_accept_ratio
  
      [
          samples_x,
          samples_y,
      ], replica_log_accept_ratio = sample.sample_chain(
          num_results=num_results,
          num_burnin_steps=200,
          current_state=[1.0, 1.0],
          kernel=remc,
          trace_fn=trace_fn,
          seed=test_util.test_seed(),
      )
      samples = tf.stack([samples_x, samples_y], axis=-1)
      sample_mean = tf.reduce_mean(samples, axis=0)
      sample_cov = sample_stats.covariance(samples, sample_axis=0)
  
      replica_accept_ratio = tf.reduce_mean(
          tf.math.exp(tf.minimum(0.0, replica_log_accept_ratio)), axis=0
      )
  
      [
          sample_mean_,
          sample_cov_,
          replica_accept_ratio_,
          ess_,
      ] = self.evaluate(
          [
              sample_mean,
              sample_cov,
              replica_accept_ratio,
              effective_sample_size(samples),
          ]
      )
      logging.vlog(
          2,
          "---- execution:{}  accept_ratio:{}  mean:{}  cov:{}".format(
              "eager" if tf.executing_eagerly() else "graph",
              replica_accept_ratio_,
              sample_mean_,
              sample_cov_,
          ),
      )
  
      self.assertAllEqual([num_results], samples_x.shape)
      self.assertAllEqual([num_results], samples_y.shape)
  
      max_scale = np.sqrt(np.max(true_cov))
  
      self.assertAllClose(
          true_mean,
          sample_mean_,
          atol=6 * max_scale / np.sqrt(np.min(ess_)),
      )
      self.assertAllClose(
          true_cov,
          sample_cov_,
          atol=6 * max_scale**2 / np.sqrt(np.min(ess_)),
      )
  
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 6 for calling the REMCTest obj's underlying member method code with example usage: obj.testMultipleCorrelatedStatesWithNoBatchDims() and returns: unknown, hmc, results