code: def _correlated_mvn_nuts(self, dim, step_size, num_steps):
      # The correlated MVN example is taken from the NUTS paper
      # https://arxiv.org/pdf/1111.4246.pdf.
      # This implementation in terms of MVNCholPrecisionTril follows
      # tfp/examples/jupyter_notebooks/Bayesian_Gaussian_Mixture_Model.ipynb
  
      class MVNCholPrecisionTriL(
          transformed_distribution.TransformedDistribution
      ):
          """MVN from loc and (Cholesky) precision matrix."""
  
          def __init__(self, loc, chol_precision_tril, name=None):
              super(MVNCholPrecisionTriL, self).__init__(
                  distribution=independent.Independent(
                      normal.Normal(
                          tf.zeros_like(loc), scale=tf.ones_like(loc)
                      ),
                      reinterpreted_batch_ndims=1,
                  ),
                  bijector=chain.Chain(
                      [
                          shift.Shift(shift=loc),
                          invert.Invert(
                              scale_matvec_tril.ScaleMatvecTriL(
                                  scale_tril=chol_precision_tril,
                                  adjoint=True,
                              )
                          ),
                      ]
                  ),
                  name=name,
              )
  
      strm = test_util.test_seed_stream()
      wishart = wishart_lib.WishartTriL(
          dim, scale_tril=tf.eye(dim), input_output_cholesky=True
      )
      chol_precision = wishart.sample(seed=strm())
      mvn = MVNCholPrecisionTriL(
          loc=tf.zeros(dim), chol_precision_tril=chol_precision
      )
      kernel = nuts_autobatching.NoUTurnSampler(
          mvn.log_prob,
          step_size=[step_size],
          num_trajectories_per_step=num_steps,
          use_auto_batching=True,
          stackless=False,
          max_tree_depth=7,
          seed=strm(),
      )
      return kernel
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import scale_matvec_tril
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.mcmc import nuts_autobatching
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 13 for calling the NutsTest obj's underlying member method code with example usage: obj._correlated_mvn_nuts(dim, step_size, num_steps) and returns: kernel