code: def assert_univariate_target_conservation(test, target_d, step_size):
      # 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
      strm = test_util.test_seed_stream()
      # 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())
      )
  
      @tf.function(autograph=False)
      def run_chain():
          kernel = nuts.NoUTurnSampler(
              target_d.log_prob,
              step_size=step_size,
              max_tree_depth=3,
              unrolled_leapfrog_steps=2,
          )
          result = sample.sample_chain(
              num_results=num_steps,
              num_burnin_steps=0,
              current_state=initialization,
              trace_fn=None,
              kernel=kernel,
              seed=strm(),
          )
          return result
  
      result = run_chain()
      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 = assert_util.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,
      )
      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 = assert_util.assert_greater_equal(
          tf.reduce_mean(movement), 1 * step_size
      )
      return (check_cdf_agrees, check_enough_power, check_movement)
  
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 2 for calling a global function with example usage: assert_univariate_target_conservation(test, target_d, step_size) and returns: check_cdf_agrees,  check_enough_power,  check_movement, result