code: def testStructuredState2(self):
      @jdab.JointDistributionCoroutineAutoBatched
      def model():
          mu = yield sample_dist_lib.Sample(
              normal.Normal(0, 1), [65], name="mu"
          )
          sigma = yield sample_dist_lib.Sample(
              exponential.Exponential(1.0), [65], name="sigma"
          )
          beta = yield sample_dist_lib.Sample(
              normal.Normal(
                  loc=tf.gather(mu, tf.range(436) % 65, axis=-1),
                  scale=tf.gather(sigma, tf.range(436) % 65, axis=-1),
              ),
              4,
              name="beta",
          )
          _ = yield multinomial.Multinomial(
              total_count=100.0,
              logits=pad.Pad([[0, 1]])(beta),
              name="y",
          )
  
      stream = test_util.test_seed_stream()
      pinned = model.experimental_pin(y=model.sample(seed=stream()).y)
      struct = pinned.dtype
      stddevs = struct._make(
          [
              tf.fill([65], 0.1),
              tf.fill([65], 1.0),
              tf.fill([436, 4], 10.0),
          ]
      )
      momentum_dist = jdab.JointDistributionNamedAutoBatched(
          struct._make(normal.Normal(0, 1 / std) for std in stddevs)
      )
      kernel = preconditioned_hmc.PreconditionedHamiltonianMonteCarlo(
          pinned.unnormalized_log_prob,
          step_size=0.1,
          num_leapfrog_steps=10,
          momentum_distribution=momentum_dist,
      )
      bijector = pinned.experimental_default_event_space_bijector()
      kernel = transformed_kernel.TransformedTransitionKernel(
          kernel, bijector
      )
      pullback_shape = bijector.inverse_event_shape(pinned.event_shape)
      kernel = dmma.DiagonalMassMatrixAdaptation(
          kernel,
          initial_running_variance=struct._make(
              sample_stats.RunningVariance.from_shape(t)
              for t in pullback_shape
          ),
      )
      state = bijector(
          struct._make(
              uniform.Uniform(-2.0, 2.0).sample(shp)
              for shp in bijector.inverse_event_shape(
                  pinned.event_shape
              )
          )
      )
      self.evaluate(
          sample_chain(
              3, current_state=state, kernel=kernel, seed=stream()
          ).all_states
      )
  
imports: import collections
  import warnings
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  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 multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import nuts
  from tensorflow_probability.python.mcmc import random_walk_metropolis
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.sample import sample_chain
  from tensorflow_probability.python.random.random_ops import rayleigh
label_desc: Use image node_id 16 for calling the SampleChainTest obj's underlying member method code with example usage: obj.testStructuredState2() without return types