code: def testSGAHMC(self):
      @tfd.JointDistributionCoroutine
      def model():
          x = yield Root(tfd.Normal(self._constant(0.0), 1.0))
          yield tfd.Sample(tfd.Normal(x, 1.0), 2)
  
      def target_log_prob_fn(x):
          return model.log_prob(x), ()
  
      @jax.jit
      def kernel(sga_hmc_state, step, seed):
          adapt = step < num_adapt_steps
          seed, hmc_seed = util.split_seed(seed, 2)
  
          (
              sga_hmc_state,
              sga_hmc_extra,
          ) = sga_hmc.stochastic_gradient_ascent_hmc_step(
              sga_hmc_state,
              scalar_step_size=self._constant(0.1),
              step_size_scale=self._constant(1.0),
              target_log_prob_fn=target_log_prob_fn,
              criterion_fn=sga_hmc.chees_criterion,
              adapt=adapt,
              seed=hmc_seed,
          )
  
          return (
              (
                  sga_hmc_state,
                  step + 1,
                  seed,
              ),
              sga_hmc_extra.trajectory_length_params.mean_trajectory_length(),
          )
  
      init_trajectory_length = self._constant(0.1)
      num_adapt_steps = 10
      _, trajectory_length = fun_mc.trace(
          (
              sga_hmc.stochastic_gradient_ascent_hmc_init(
                  util.map_tree_up_to(
                      model.dtype,
                      lambda dtype, shape: jnp.zeros(  # pylint: disable=g-long-lambda
                          (16,) + tuple(shape), dtype
                      ),
                      model.dtype,
                      model.event_shape,
                  ),
                  target_log_prob_fn,
                  init_trajectory_length=init_trajectory_length,
              ),
              0,
              self._make_seed(_test_seed()),
          ),
          kernel,
          num_adapt_steps + 2,
      )
  
      # We expect it to increase as part of adaptation.
      self.assertAllGreater(
          trajectory_length[-1], init_trajectory_length
      )
      # After adaptation is done, the trajectory length should remain constant.
      self.assertAllClose(trajectory_length[-1], trajectory_length[-2])
  
imports: import functools
  import os
  from absl.testing import parameterized
  from jax import config
  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 sga_hmc
  from fun_mc import test_util
label_desc: Use image node_id 6 for calling the SGAHMCTest obj's underlying member method code with example usage: obj.testSGAHMC() and returns: model,  , sga_hmc_state,  unknown,  seed,  sga_hmc_extra