code: def testAdaptiveHMC(self):
      num_chains = 16
      num_steps = 4000
      num_warmup_steps = num_steps // 2
      num_adapt_steps = int(0.8 * num_warmup_steps)
  
      # Setup the model and state constraints.
      model = tfp.distributions.JointDistributionSequential(
          [
              tfp.distributions.Normal(
                  loc=self._constant(0.0), scale=1.0
              ),
              tfp.distributions.Independent(
                  tfp.distributions.LogNormal(
                      loc=self._constant([1.0, 1.0]), scale=0.5
                  ),
                  1,
              ),
          ]
      )
      bijector = [tfp.bijectors.Identity(), tfp.bijectors.Exp()]
      transform_fn = util_tfp.bijector_to_transform_fn(
          bijector, model.dtype, batch_ndims=1
      )
  
      def target_log_prob_fn(*x):
          return model.log_prob(x), ()
  
      # Start out at zeros (in the unconstrained space).
      state, _ = transform_fn(
          *[
              jnp.zeros([num_chains] + list(e), dtype=self._dtype)
              for e in model.event_shape
          ]
      )
  
      (
          reparam_log_prob_fn,
          reparam_state,
      ) = fun_mc.reparameterize_potential_fn(
          target_log_prob_fn, transform_fn, state
      )
  
      # Define the kernel.
      def kernel(adaptive_hmc_state, seed):
          hmc_seed, seed = util.split_seed(seed, 2)
  
          (
              adaptive_hmc_state,
              adaptive_hmc_extra,
          ) = prefab.adaptive_hamiltonian_monte_carlo_step(
              adaptive_hmc_state,
              target_log_prob_fn=reparam_log_prob_fn,
              num_adaptation_steps=num_adapt_steps,
              seed=hmc_seed,
          )
  
          return (adaptive_hmc_state, seed), (
              adaptive_hmc_extra.state,
              adaptive_hmc_extra.is_accepted,
              adaptive_hmc_extra.step_size,
          )
  
      seed = self._make_seed(_test_seed())
  
      _, (state_chain, is_accepted_chain, _) = jax.jit(
          lambda reparam_state, seed: fun_mc.trace(  # pylint: disable=g-long-lambda
              state=(
                  prefab.adaptive_hamiltonian_monte_carlo_init(
                      reparam_state, reparam_log_prob_fn
                  ),
                  seed,
              ),
              fn=kernel,
              num_steps=num_steps,
          )
      )(reparam_state, seed)
  
      # Discard the warmup samples.
      state_chain = [s[num_warmup_steps:] for s in state_chain]
      is_accepted_chain = is_accepted_chain[num_warmup_steps:]
  
      accept_rate = jnp.mean(
          jnp.asarray(is_accepted_chain, jnp.float32)
      )
      rhat = tfp.mcmc.potential_scale_reduction(state_chain)
      sample_mean = [jnp.mean(s, axis=[0, 1]) for s in state_chain]
      sample_var = [jnp.var(s, axis=[0, 1]) for s in state_chain]
  
      self.assertAllAssertsNested(
          lambda rhat: self.assertAllLess(rhat, 1.1), rhat
      )
      self.assertAllClose(0.8, accept_rate, atol=0.05)
      self.assertAllClose(model.mean(), sample_mean, rtol=0.1, atol=0.1)
      self.assertAllClose(
          model.variance(), sample_var, rtol=0.1, atol=0.1
      )
  
imports: import functools
  import os
  import jax
  from jax import config
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import prefab
  from fun_mc import test_util
  from fun_mc import util_tfp
label_desc: Use image node_id 4 for calling the PrefabTest obj's underlying member method code with example usage: obj.testAdaptiveHMC() and returns: model,  , adaptive_hmc_state,  seed,  adaptive_hmc_extra,  adaptive_hmc_extra,  adaptive_hmc_extra