code: def run_nuts_chain(
      event_size,
      batch_size,
      num_steps,
      initial_state=None,
      dry_run=False,
      stackless=False,
      seed=None,
  ):
      if seed is None:
          seed = 1
  
      def target_log_prob_fn(event):
          with tf1.name_scope(
              "nuts_test_target_log_prob", values=[event]
          ):
              return mvn_diag.MultivariateNormalDiag(
                  tf.zeros(event_size)
              ).log_prob(event)
  
      if initial_state is not None:
          state = initial_state
      else:
          state = tf.zeros([batch_size, event_size])
      # kernel = tfp.mcmc.HamiltonianMonteCarlo(
      #     target_log_prob_fn,
      #     num_leapfrog_steps=3,
      #     step_size=0.3,
      #     seed=seed)
      kernel = nuts_autobatching.NoUTurnSampler(
          target_log_prob_fn,
          step_size=[0.3],
          use_auto_batching=not dry_run,
          stackless=stackless,
          unrolled_leapfrog_steps=2,
          max_tree_depth=4,
          seed=seed,
      )
      chain_state, extra = 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=[state],
          kernel=kernel,
          parallel_iterations=1,
      )
      return chain_state, extra.leapfrogs_taken
  
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 1 for calling a global function with example usage: run_nuts_chain(event_size, batch_size, num_steps, initial_state, dry_run, stackless, seed) and returns: chain_state,  extra, mvn_diag