code: def testStepSizeAdaptation(self):
      def log_accept_ratio_fn(step_size):
          return -(step_size**2)
  
      def kernel(ssa_state, seed):
          normal_seed, seed = util.split_seed(seed, 2)
          log_accept_ratio = log_accept_ratio_fn(
              ssa_state.step_size()
          ) + 0.01 * util.random_normal([4], self._dtype, normal_seed)
          ssa_state, ssa_extra = prefab.step_size_adaptation_step(
              ssa_state, log_accept_ratio, num_adaptation_steps=100
          )
          return (ssa_state, seed), (
              ssa_extra.accept_prob,
              ssa_state.step_size(),
              ssa_state.step_size(num_adaptation_steps=100),
          )
  
      seed = self._make_seed(_test_seed())
  
      _, (p_accept, step_size, rms_step_size) = fun_mc.trace(
          (
              prefab.step_size_adaptation_init(
                  jnp.array(0.1, self._dtype)
              ),
              seed,
          ),
          kernel,
          200,
      )
  
      self.assertAllClose(0.8, p_accept[100], atol=0.1)
      self.assertAllClose(step_size[100], step_size[150])
      self.assertAllClose(rms_step_size[100], rms_step_size[150])
  
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 6 for calling the PrefabTest obj's underlying member method code with example usage: obj.testStepSizeAdaptation() and returns: unknown, ssa_state,  seed,  ssa_extra,  ssa_state,  ssa_state