code: def make_test_nuts_kernel_fn(
      target_log_prob_fn, init_state, scalings
  ):
      """Set up a function to generate nuts kernel for testing."""
      max_tree_depth = 3
  
      state_std = [
          tf.math.reduce_std(x, axis=0, keepdims=True)
          for x in init_state
      ]
      step_size = compute_hmc_step_size(
          scalings, state_std, max_tree_depth**2
      )
      return nuts.NoUTurnSampler(
          target_log_prob_fn=target_log_prob_fn,
          step_size=step_size,
          max_tree_depth=max_tree_depth,
      )
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mixture
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.distributions.internal import statistical_testing
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import compute_hmc_step_size
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import gen_make_hmc_kernel_fn
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import gen_make_transform_hmc_kernel_fn
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import make_rwmh_kernel_fn
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import sample_sequential_monte_carlo
  from tensorflow_probability.python.experimental.mcmc.sample_sequential_monte_carlo import simple_heuristic_tuning
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import nuts
label_desc: Use image node_id 1 for calling a global function with example usage: make_test_nuts_kernel_fn(target_log_prob_fn, init_state, scalings) and returns: nuts