code: def testDocstringExample(self):
      if tf.executing_eagerly() and not JAX_MODE:
          self.skipTest("Too slow for TF Eager.")
  
      target = jds.JointDistributionSequential(
          [
              normal.Normal(0.0, tf.constant(20.0, dtype=self.dtype)),
              half_normal.HalfNormal(
                  tf.constant(10.0, dtype=self.dtype)
              ),
          ]
      )
  
      def target_log_prob_fn(*x):
          return tf.cast(target.log_prob(x), self.dtype)
  
      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_fn,
          step_size=step_size,
          num_leapfrog_steps=1,
      )
      kernel = gbtla.GradientBasedTrajectoryLengthAdaptation(
          kernel,
          num_adaptation_steps=num_adaptation_steps,
          validate_args=True,
      )
      kernel = dassa.DualAveragingStepSizeAdaptation(
          kernel, num_adaptation_steps=num_adaptation_steps
      )
      kernel = transformed_kernel.TransformedTransitionKernel(
          kernel, [identity.Identity(), exp.Exp()]
      )
  
      def trace_fn(_, pkr):
          return (
              pkr.inner_results.inner_results.inner_results.accepted_results.step_size,
              pkr.inner_results.inner_results.max_trajectory_length,
              pkr.inner_results.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.ones(num_chains, dtype=self.dtype),
              tf.ones(num_chains, 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(0.52, mean_step_size, atol=0.2)
      self.assertAllClose(46.0, mean_max_trajectory_length, atol=15)
      self.assertAllClose(
          target.mean(),
          [tf.reduce_mean(x, axis=[0, 1]) for x in chain],
          atol=1.5,
      )
      self.assertAllClose(
          target.variance(),
          [tf.math.reduce_variance(x, axis=[0, 1]) for x in chain],
          rtol=0.2,
      )
  
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 1 for calling the _GradientBasedTrajectoryLengthAdaptationTest obj's underlying member method code with example usage: obj.testDocstringExample() and returns: tf, pkr,  pkr,  pkr