code: def test_xla_compile_and_correctness(self):
      tf.set_random_seed(3)
      nsamples = 1000
      nchains = 10
      nd = 5
      theta0 = np.zeros((nchains, nd))
      mu = np.arange(nd)
      stddev = np.arange(nd) + 1.0
      rng = np.random.RandomState(seed=4)
      step_size = rng.rand(nchains, 1) * 0.1 + 1
  
      num_steps = nsamples
      initial_state = tf.cast(theta0, dtype=tf.float32)
      unrolled_leapfrog_steps = 2
  
      @tf.function(autograph=False)
      def target_log_prob_fn(event):
          with tf.name_scope("nuts_test_target_log_prob"):
              return mvn_diag.MultivariateNormalDiag(
                  loc=tf.cast(mu, dtype=tf.float32),
                  scale_diag=tf.cast(stddev, dtype=tf.float32),
              ).log_prob(event)
  
      @tf.function(autograph=False, jit_compile=True)
      def _run_nuts_chain():
          kernel = nuts.NoUTurnSampler(
              target_log_prob_fn,
              step_size=[tf.cast(step_size, dtype=tf.float32)],
              unrolled_leapfrog_steps=unrolled_leapfrog_steps,
          )
          [x], (is_accepted, leapfrogs_taken) = sample.sample_chain(
              num_results=num_steps,
              num_burnin_steps=0,
              current_state=[initial_state],
              kernel=kernel,
              trace_fn=lambda _, pkr: (
                  pkr.is_accepted,
                  pkr.leapfrogs_taken,
              ),
              seed=9,
          )
          return (
              tf.reduce_mean(x, axis=[0, 1]),
              tf.math.reduce_std(x, axis=[0, 1]),
              is_accepted,
              leapfrogs_taken,
          )
  
      with tf.device(FLAGS.test_device):
          [
              sample_mean,
              sample_stddev,
              is_accepted,
              leapfrogs_taken,
          ] = self.evaluate(_run_nuts_chain())
  
      self.assertAllClose(mu, sample_mean, atol=0.1, rtol=0.1)
      self.assertAllClose(stddev, sample_stddev, atol=0.15, rtol=0.15)
      # Test early stopping in tree building
      self.assertTrue(
          np.any(
              np.isin(
                  np.asarray([5, 9, 11, 13]) * unrolled_leapfrog_steps,
                  np.unique(leapfrogs_taken[is_accepted]),
              )
          )
      )
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/nuts_autobatching_xla_test.py
imports: from absl import flags
  import numpy
  import tensorflow.compat.v1
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.experimental.auto_batching import tf_backend
  from tensorflow_probability.python.experimental.mcmc import nuts_autobatching
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 2 for calling the NutsXLATest obj's underlying member method code with example usage: obj.test_xla_compile_and_correctness() and returns: tf,  tf,  is_accepted,  leapfrogs_taken, mvn_diag