code: def assert_univariate_target_conservation(
      test, mk_target, step_size, stackless
  ):
      # Sample count limited partly by memory reliably available on Forge.  The test
      # remains reasonable even if the nuts recursion limit is severely curtailed
      # (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump
      # the sample count.
      num_samples = int(5e4)
      num_steps = 1
      target_d = mk_target()
      strm = test_util.test_seed_stream(
          salt="univariate_nuts_test", hardcoded_seed=1
      )
      # We wrap the initial values in `tf.identity` to avoid broken gradients
      # resulting from a bijector cache hit, since bijectors of the same
      # type/parameterization now share a cache.
      # TODO(b/72831017): Fix broken gradients caused by bijector caching.
      initialization = tf.identity(
          target_d.sample([num_samples], seed=strm())
      )
  
      def target(*args):
          # TODO(axch): Just use target_d.log_prob directly, and accept target_d
          # itself as an argument instead of a maker function.  Blocked by
          # b/128932888.  It would then also be nice not to eta-expand
          # target_d.log_prob; that was blocked by b/122414321, but maybe tfp's port
          # of value_and_gradients_function fixed that bug.
          return mk_target().log_prob(*args)
  
      operator = nuts_autobatching.NoUTurnSampler(
          target,
          step_size=step_size,
          max_tree_depth=3,
          use_auto_batching=True,
          stackless=stackless,
          unrolled_leapfrog_steps=2,
          seed=strm(),
      )
      result, extra = sample.sample_chain(
          num_results=num_steps,
          num_burnin_steps=0,
          current_state=initialization,
          kernel=operator,
      )
      # Note: sample_chain puts the chain history on top, not the (independent)
      # chains.
      test.assertAllEqual([num_steps, num_samples], result.shape)
      answer = result[0]
      check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(
          answer, target_d.cdf, false_fail_rate=1e-6
      )
      check_enough_power = tf1.assert_less(
          st.min_discrepancy_of_true_cdfs_detectable_by_dkwm(
              num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6
          ),
          0.025,
      )
      test.assertAllEqual([num_samples], extra.leapfrogs_taken[0].shape)
      unique, _ = tf.unique(extra.leapfrogs_taken[0])
      check_leapfrogs_vary = tf1.assert_greater_equal(
          tf.shape(input=unique)[0], 3
      )
      avg_leapfrogs = tf.math.reduce_mean(
          input_tensor=extra.leapfrogs_taken[0]
      )
      check_leapfrogs = tf1.assert_greater_equal(
          avg_leapfrogs, tf.constant(4, dtype=avg_leapfrogs.dtype)
      )
      movement = tf.abs(answer - initialization)
      test.assertAllEqual([num_samples], movement.shape)
      # This movement distance (1 * step_size) was selected by reducing until 100
      # runs with independent seeds all passed.
      check_movement = tf1.assert_greater_equal(
          tf.reduce_mean(input_tensor=movement), 1 * step_size
      )
      return (
          check_cdf_agrees,
          check_enough_power,
          check_leapfrogs_vary,
          check_leapfrogs,
          check_movement,
      )
  
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 2 for calling a global function with example usage: assert_univariate_target_conservation(test, mk_target, step_size, stackless) and returns: check_cdf_agrees,  check_enough_power,  check_leapfrogs_vary,  check_leapfrogs,  check_movement, mk_target