code: def test_supply_full_step_size(self):
      if tf.executing_eagerly() and not JAX_MODE:
          self.skipTest("Eager test is slow.")
      stream = test_util.test_seed_stream()
  
      jd_model = jdn.JointDistributionNamed(
          {
              "a": normal.Normal(0.0, 1.0),
              "b": mvn_diag.MultivariateNormalDiag(
                  loc=tf.zeros(3),
                  scale_diag=tf.constant([1.0, 2.0, 3.0]),
              ),
          }
      )
  
      init_step_size = {
          "a": tf.reshape(tf.linspace(1.0, 2.0, 3), (3, 1)),
          "b": tf.reshape(tf.linspace(1.0, 2.0, 9), (3, 3)),
      }
  
      _, actual_step_size = windowed_sampling.windowed_adaptive_hmc(
          1,
          jd_model,
          num_adaptation_steps=25,
          n_chains=3,
          init_step_size=init_step_size,
          num_leapfrog_steps=5,
          discard_tuning=False,
          trace_fn=lambda *args: unnest.get_innermost(
              args[-1], "step_size"
          ),
          seed=stream(),
      )
  
      # Gets a newaxis because step size needs to have an event dimension.
      self.assertAllCloseNested(
          [init_step_size["a"], init_step_size["b"]],
          [j[0] for j in actual_step_size],
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.distributions import autoregressive
  from tensorflow_probability.python.distributions import beta_binomial
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental import distribute
  from tensorflow_probability.python.experimental.mcmc import windowed_sampling
  from tensorflow_probability.python.internal import callable_util
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import diagnostic
label_desc: Use image node_id 1 for calling the WindowedSamplingStepSizeTest obj's underlying member method code with example usage: obj.test_supply_full_step_size() without return types