code: def test_adaptation(self, criterion_fn):
      # Compare this to testAdaptation. There we don't use SPMD, but should
      # get the same hyperparameters.
  
      if not JAX_MODE:
          self.skipTest("TF does not have pmax implemented.")
  
      target = independent.Independent(
          normal.Normal(0.0, tf.constant([1.0, 10.0])), 1
      )
  
      def run(seed):
          num_burnin_steps = 1000
          num_adaptation_steps = int(num_burnin_steps * 0.8)
          num_results = 500
          num_chains = 16 // distribute_test_lib.NUM_DEVICES
          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,
              experimental_reduce_chain_axis_names=self.axis_name,
              validate_args=True,
          )
          kernel = dassa.DualAveragingStepSizeAdaptation(
              kernel,
              num_adaptation_steps=num_adaptation_steps,
              experimental_reduce_chain_axis_names=self.axis_name,
          )
  
          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]),
              kernel=kernel,
              trace_fn=trace_fn,
              seed=seed,
          )
  
          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
          )
          mean = tf.reduce_mean(chain, axis=[0, 1])
          var = tf.reduce_variance(chain, axis=[0, 1])
  
          return (
              mean,
              var,
              p_accept,
              mean_step_size,
              mean_max_trajectory_length,
          )
  
      seeds = self.shard_values(
          tf.stack(
              samplers.split_seed(
                  samplers.zeros_seed(), distribute_test_lib.NUM_DEVICES
              )
          ),
          0,
      )
  
      (
          mean,
          var,
          p_accept,
          mean_step_size,
          mean_max_trajectory_length,
      ) = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(
                  run, args=(seeds,), axis_name=self.axis_name
              ),
              0,
          )
      )
  
      self.assertAllClose(0.75, p_accept.mean(), atol=0.1)
      # Both ChEES-rate and SNAPER learn roughly the same trajectory length.
      self.assertAllClose(1.5, mean_step_size[0], atol=0.2)
      self.assertAllClose(15.0, mean_max_trajectory_length[0], rtol=0.3)
      self.assertAllClose(target.mean(), mean.mean(0), atol=1.0)
      self.assertAllClose(
          target.variance(), var.mean(0) + mean.var(0), 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 4 for calling the DistributedGBTLATest obj's underlying member method code with example usage: obj.test_adaptation(criterion_fn) and returns: mean,  var,  p_accept,  mean_step_size,  mean_max_trajectory_length, pkr,  pkr,  pkr