code: def testMultivariateNormalNdConvergence(self, nsamples, nchains, nd):
      strm = test_util.test_seed_stream()
      theta0 = np.zeros((nchains, nd))
      mu = np.arange(nd)
      w = np.random.randn(nd, nd) * 0.1
      cov = w * w.T + np.diagflat(np.arange(nd) + 1.0)
      step_size = np.random.rand(nchains, 1) * 0.1 + 1.0
  
      @tf.function(autograph=False)
      def run_chain_and_get_summary(
          mu, scale_tril, step_size, nsamples, state
      ):
          def target_log_prob_fn(event):
              with tf.name_scope("nuts_test_target_log_prob"):
                  return mvn_tril.MultivariateNormalTriL(
                      loc=tf.cast(mu, dtype=tf.float64),
                      scale_tril=tf.cast(scale_tril, dtype=tf.float64),
                  ).log_prob(event)
  
          kernel = nuts.NoUTurnSampler(
              target_log_prob_fn,
              step_size=[step_size],
              max_tree_depth=5,
          )
  
          def trace_fn(_, pkr):
              return (pkr.is_accepted, pkr.leapfrogs_taken)
  
          [x], (is_accepted, leapfrogs_taken) = sample.sample_chain(
              num_results=nsamples,
              num_burnin_steps=0,
              current_state=[tf.cast(state, dtype=tf.float64)],
              kernel=kernel,
              trace_fn=trace_fn,
              seed=strm(),
          )
  
          return (
              tf.shape(x),
              # We'll average over samples (dim=0) and chains (dim=1).
              tf.reduce_mean(x, axis=[0, 1]),
              sample_stats.covariance(x, sample_axis=[0, 1]),
              leapfrogs_taken[is_accepted],
          )
  
      (
          sample_shape,
          sample_mean,
          sample_cov,
          leapfrogs_taken,
      ) = self.evaluate(
          run_chain_and_get_summary(
              mu, np.linalg.cholesky(cov), step_size, nsamples, theta0
          )
      )
  
      self.assertAllEqual(sample_shape, [nsamples, nchains, nd])
      self.assertAllClose(mu, sample_mean, atol=0.1, rtol=0.1)
      self.assertAllClose(cov, sample_cov, atol=0.15, rtol=0.15)
      # Test early stopping in tree building
      self.assertTrue(
          np.any(
              np.isin(
                  np.asarray([5, 9, 11, 13]), np.unique(leapfrogs_taken)
              )
          )
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import sigmoid_beta
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 9 for calling the NutsTest obj's underlying member method code with example usage: obj.testMultivariateNormalNdConvergence(nsamples, nchains, nd) and returns: tf,  tf,  sample_stats,  leapfrogs_taken, pkr,  pkr, mvn_tril