code: def testAdaptation(self, criterion_fn):
      if tf.executing_eagerly() and not JAX_MODE:
          self.skipTest("Too slow for TF Eager.")
  
      target = independent.Independent(
          normal.Normal(0.0, tf.constant([1.0, 10.0], self.dtype)), 1
      )
  
      num_burnin_steps = 1000
      num_adaptation_steps = int(num_burnin_steps * 0.8)
      num_results = 500
      num_chains = 16
      step_size = 0.1
  
      kernel = hmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target.log_prob,
          step_size=step_size,
          num_leapfrog_steps=1,
      )
      kernel = gbtla.GradientBasedTrajectoryLengthAdaptation(
          kernel,
          num_adaptation_steps=num_adaptation_steps,
          criterion_fn=criterion_fn,
          validate_args=True,
      )
      kernel = dassa.DualAveragingStepSizeAdaptation(
          kernel, num_adaptation_steps=num_adaptation_steps
      )
  
      def trace_fn(_, pkr):
          return (
              pkr.inner_results.inner_results.accepted_results.step_size,
              pkr.inner_results.max_trajectory_length,
              pkr.inner_results.inner_results.log_accept_ratio,
          )
  
      # The chain will be stepped for num_results + num_burnin_steps, adapting for
      # the first num_adaptation_steps.
      chain, [
          step_size,
          max_trajectory_length,
          log_accept_ratio,
      ] = sample.sample_chain(
          num_results=num_results,
          num_burnin_steps=num_burnin_steps,
          current_state=tf.zeros([num_chains, 2], dtype=self.dtype),
          kernel=kernel,
          trace_fn=trace_fn,
          seed=test_util.test_seed(sampler_type="stateless"),
      )
  
      p_accept = tf.math.exp(
          generic.reduce_logmeanexp(tf.minimum(log_accept_ratio, 0.0))
      )
      mean_step_size = tf.reduce_mean(step_size)
      mean_max_trajectory_length = tf.reduce_mean(max_trajectory_length)
  
      self.assertAllClose(0.75, p_accept, atol=0.1)
      self.assertAllClose(1.5, mean_step_size, atol=0.2)
      # Both SNAPER and ChEES-rate find roughly the same trajectory length for
      # this target.
      self.assertAllClose(15.0, mean_max_trajectory_length, rtol=0.3)
      self.assertAllClose(
          target.mean(), tf.reduce_mean(chain, axis=[0, 1]), atol=1.0
      )
      self.assertAllClose(
          target.variance(),
          tf.math.reduce_variance(chain, axis=[0, 1]),
          rtol=0.1,
      )
  
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 half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.mcmc import gradient_based_trajectory_length_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  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 dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
label_desc: Use image node_id 6 for calling the _GradientBasedTrajectoryLengthAdaptationTest obj's underlying member method code with example usage: obj.testAdaptation(criterion_fn) and returns: pkr,  pkr,  pkr