code: def run_nuts_chain(
      event_size, batch_size, num_steps, initial_state=None, seed=None
  ):
      if seed is None:
          seed = test_util.test_seed()
  
      def target_log_prob_fn(event):
          with tf.name_scope("nuts_test_target_log_prob"):
              return mvn_diag.MultivariateNormalDiag(
                  tf.zeros(event_size)
              ).log_prob(event)
  
      if initial_state is None:
          initial_state = tf.zeros([batch_size, event_size])
  
      kernel = nuts.NoUTurnSampler(
          target_log_prob_fn,
          step_size=[0.3],
          unrolled_leapfrog_steps=2,
          max_tree_depth=4,
      )
  
      chain_state, leapfrogs_taken = sample.sample_chain(
          num_results=num_steps,
          num_burnin_steps=0,
          # Intentionally pass a list argument to test that singleton lists are
          # handled reasonably (c.f. assert_univariate_target_conservation, which
          # uses an unwrapped singleton).
          current_state=[initial_state],
          kernel=kernel,
          trace_fn=lambda _, pkr: pkr.leapfrogs_taken,
          seed=seed,
      )
  
      return chain_state, 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 1 for calling a global function with example usage: run_nuts_chain(event_size, batch_size, num_steps, initial_state, seed) and returns: chain_state,  leapfrogs_taken, mvn_diag