code: def testChainWorksCorrelatedMultivariate(self):
      # Disabling this test in eager mode as it is very slow.
      # Other tests check eager behavior, so we can safely disable this test.
      if tf.executing_eagerly() and not JAX_MODE:
          return
  
      dtype = np.float32
      true_mean = dtype([0, 0])
      true_cov = dtype([[1, 0.5], [0.5, 1]])
      true_cov_chol = np.linalg.cholesky(true_cov)
      num_results = 3000
      counter = collections.Counter()
  
      def target_log_prob(x, y):
          counter["target_calls"] += 1
          # Corresponds to unnormalized MVN.
          # z = matmul(inv(chol(true_cov)), [x, y] - true_mean)
          z = tf.stack([x, y], axis=-1) - true_mean
          z = tf.linalg.triangular_solve(
              true_cov_chol, z[..., tf.newaxis]
          )[..., 0]
          return -0.5 * tf.reduce_sum(z**2.0, axis=-1)
  
      states = sample_chain(
          num_results=num_results,
          current_state=[dtype(-2), dtype(2)],
          kernel=hmc.HamiltonianMonteCarlo(
              target_log_prob_fn=target_log_prob,
              step_size=[0.5, 0.5],
              num_leapfrog_steps=2,
          ),
          num_burnin_steps=200,
          num_steps_between_results=1,
          trace_fn=None,
          seed=test_util.test_seed(),
      )
      self.assertAllEqual(dict(target_calls=2), counter)
      states = tf.stack(states, axis=-1)
      self.assertEqual(
          num_results, tf.compat.dimension_value(states.shape[0])
      )
      sample_mean = tf.reduce_mean(states, axis=0)
      x = states - sample_mean
      sample_cov = tf.matmul(x, x, transpose_a=True) / dtype(
          num_results
      )
      sample_mean_, sample_cov_ = self.evaluate(
          [sample_mean, sample_cov]
      )
      self.assertAllClose(true_mean, sample_mean_, atol=0.1, rtol=0.0)
      self.assertAllClose(true_cov, sample_cov_, atol=0.0, rtol=0.175)
  
imports: import collections
  import warnings
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import random_walk_metropolis
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.sample import sample_chain
  from tensorflow_probability.python.random.random_ops import rayleigh
label_desc: Use image node_id 2 for calling the SampleChainTest obj's underlying member method code with example usage: obj.testChainWorksCorrelatedMultivariate() and returns: unknown