code: def testStructuredState(
      self, kernel_from_log_prob, skip="", **model_kwargs
  ):
      if skip:
          self.skipTest(skip)
      seed_stream = test_util.test_seed_stream()
  
      n = 300
      p = 50
      x = tf.random.normal([n, p], seed=seed_stream())
  
      def beta_proportion(mu, kappa):
          return beta_lib.Beta(
              concentration0=mu * kappa, concentration1=(1 - mu) * kappa
          )
  
      root = jdc.JointDistributionCoroutine.Root
  
      def model_coroutine():
          beta = yield root(
              sample_dist_lib.Sample(
                  normal.Normal(0, 1), [p], name="beta"
              )
          )
          alpha = yield root(normal.Normal(0, 1, name="alpha"))
          kappa = yield root(gamma.Gamma(1, 1, name="kappa"))
          mu = tf.math.sigmoid(
              alpha[..., tf.newaxis]
              + tf.einsum("...p,np->...n", beta, x)
          )
          yield independent.Independent(
              beta_proportion(mu, kappa[..., tf.newaxis]),
              reinterpreted_batch_ndims=1,
              name="prob",
          )
  
      model = jdc.JointDistributionCoroutine(
          model_coroutine, **model_kwargs
      )
      probs = model.sample(seed=seed_stream())[-1]
      pinned = model.experimental_pin(prob=probs)
  
      kernel = kernel_from_log_prob(pinned.unnormalized_log_prob)
      nburnin = 5
      if not isinstance(
          kernel, random_walk_metropolis.RandomWalkMetropolis
      ):
          kernel = ssa.SimpleStepSizeAdaptation(
              kernel, num_adaptation_steps=nburnin // 2
          )
      kernel = transformed_kernel.TransformedTransitionKernel(
          kernel, pinned.experimental_default_event_space_bijector()
      )
      nchains = 4
  
      @tf.function
      def sample():
          return sample_chain(
              1,
              current_state=pinned.sample_unpinned(
                  nchains, seed=seed_stream()
              ),
              kernel=kernel,
              num_burnin_steps=nburnin,
              trace_fn=None,
              seed=seed_stream(),
          )
  
      self.evaluate(sample())
  
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 15 for calling the SampleChainTest obj's underlying member method code with example usage: obj.testStructuredState(kernel_from_log_prob, skip) and returns: beta_lib, sample_chain