code: def test_batches_of_chains(self, kind, n_chains):
      def model_fn():
          x = yield mvn_diag.MultivariateNormalDiag(
              tf.zeros(3), tf.ones(3), name="x"
          )
          yield multinomial.Multinomial(
              logits=pad.Pad([(0, 1)])(x), total_count=10, name="y"
          )
  
      model = jdab.JointDistributionCoroutineAutoBatched(
          model_fn, batch_ndims=1
      )
      samp = model.sample(seed=test_util.test_seed())
      states, trace = self.evaluate(
          windowed_sampling.windowed_adaptive_hmc(
              5,
              model.experimental_pin(y=samp.y),
              n_chains=n_chains,
              num_leapfrog_steps=3,
              num_adaptation_steps=25,
              seed=test_util.test_seed(),
          )
      )
      if isinstance(n_chains, int):
          n_chains = [n_chains]
      self.assertEqual((5, *n_chains, 3), states.x.shape)
      self.assertEqual((5,), trace["step_size"].shape)
  
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 16 for calling the WindowedSamplingTest obj's underlying member method code with example usage: obj.test_batches_of_chains(kind, n_chains) without return types