class_docstr: Use gradient ascent to adapt inner kernel's trajectory length.
  
  This kernel optimizes the continuous trajectory length (aka integration time)
  parameter of Hamiltonian Monte Carlo. It does so by following the gradient of
  a criterion with respect to the trajectory length. The criterion is computed
  via `criterion_fn` with signature `(previous_state, proposed_state,
  accept_prob, trajectory_length) -> criterion`, where both the returned
  values retain the batch dimensions implied by the first three inputs. See
  `chees_criterion` for an example.
  
  To avoid resonances, this kernel jitters the integration time between 0 and
  the learned trajectory length by default.
  
  The initial trajectory length is extracted from the inner
  `HamiltonianMonteCarlo` kernel by multiplying the initial step size and
  initial number of leapfrog steps. This (and other algorithmic details) imply
  that the step size must be a scalar.
  
  In general, adaptation prevents the chain from reaching a stationary
  distribution, so obtaining consistent samples requires `num_adaptation_steps`
  be set to a value [somewhat smaller][1] than the number of burnin steps.
  However, it may sometimes be helpful to set `num_adaptation_steps` to a larger
  value during development in order to inspect the behavior of the chain during
  adaptation.
  
  #### Examples
  
  This implements something similar to ChEES HMC from [2].
  
  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfb = tfp.bijectors
  tfd = tfp.distributions
  
  target_log_prob_fn = tfd.JointDistributionSequential([
      tfd.Normal(0., 20.),
      tfd.HalfNormal(10.),
  ]).log_prob
  
  num_burnin_steps = 1000
  num_adaptation_steps = int(num_burnin_steps * 0.8)
  num_results = 500
  num_chains = 16
  step_size = 0.1
  
  kernel = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=target_log_prob_fn,
      step_size=step_size,
      num_leapfrog_steps=1,
  )
  kernel = tfp.experimental.mcmc.GradientBasedTrajectoryLengthAdaptation(
      kernel,
      num_adaptation_steps=num_adaptation_steps)
  kernel = tfp.mcmc.DualAveragingStepSizeAdaptation(
      kernel, num_adaptation_steps=num_adaptation_steps)
  kernel = tfp.mcmc.TransformedTransitionKernel(
      kernel,
      [tfb.Identity(),
       tfb.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.
  samples, [step_size, max_trajectory_length, log_accept_ratio] = (
      tfp.mcmc.sample_chain(
          num_results=num_results,
          num_burnin_steps=num_burnin_steps,
          current_state=[tf.zeros(num_chains),
                         tf.zeros(num_chains)],
          kernel=kernel,
          trace_fn=trace_fn,))
  
  # ~0.75
  accept_prob = tf.math.exp(tfp.math.reduce_logmeanexp(
      tf.minimum(log_accept_ratio, 0.)))
  ```
  
  #### References
  
  [1]: <http://andrewgelman.com/2017/12/15/
        burn-vs-warm-iterative-simulation-algorithms/#comment-627745>
  
  [2]: Hoffman, M., Radul, A., & Sountsov, P. (2020). An Adaptive MCMC Scheme
       for Setting Trajectory Lengths in Hamiltonian Monte Carlo.
       <https://proceedings.mlr.press/v130/hoffman21a>
code: def proposed_state_getter_fn(self, kernel_results):
      return self._parameters["proposed_state_getter_fn"](
          kernel_results
      )
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import nest
label_desc: Use image node_id 13 for calling the GradientBasedTrajectoryLengthAdaptation obj's underlying member method code with example usage: obj.proposed_state_getter_fn(kernel_results) and returns: self