code: def run_nuts_chain(event_size, batch_size, num_steps):
      def f():
          def target_log_prob_fn(event):
              return mvn_diag.MultivariateNormalDiag(
                  tf.zeros(event_size)
              ).log_prob(event)
  
          state = tf.zeros([batch_size, event_size])
          chain_state, extra = sample.sample_chain(
              num_results=num_steps,
              num_burnin_steps=0,
              current_state=[state],
              kernel=nuts_exp.NoUTurnSampler(
                  target_log_prob_fn,
                  step_size=[0.3],
                  use_auto_batching=True,
                  seed=1,
                  backend=tf_backend.TensorFlowBackend(
                      safety_checks=False, while_parallel_iterations=1
                  ),
              ),
              parallel_iterations=1,
          )
          return chain_state, extra.leapfrogs_taken
  
      return f
  
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 1 for calling a global function with example usage: run_nuts_chain(event_size, batch_size, num_steps) and returns: f, chain_state,  extra, mvn_diag