code: def testMultipleStateParts(self):
      dist = jds.JointDistributionSequential(
          [
              mvn_diag.MultivariateNormalDiag(tf.zeros(3), tf.ones(3)),
              mvn_diag.MultivariateNormalDiag(tf.zeros(2), tf.ones(2)),
          ]
      )
      target_log_prob_fn = dist.log_prob
      kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          num_leapfrog_steps=2,
          step_size=1.0,
      )
      initial_running_variance = [
          sample_stats.RunningVariance.from_stats(
              num_samples=1.0, mean=tf.zeros(3), variance=tf.ones(3)
          ),
          sample_stats.RunningVariance.from_stats(
              num_samples=1.0, mean=tf.zeros(2), variance=tf.ones(2)
          ),
      ]
      kernel = dmma.DiagonalMassMatrixAdaptation(
          inner_kernel=kernel,
          initial_running_variance=initial_running_variance,
      )
  
      num_results = 5
      draws = sample.sample_chain(
          num_results=num_results,
          current_state=[tf.zeros(3), tf.zeros(2)],
          kernel=kernel,
          seed=test_util.test_seed(),
          trace_fn=None,
      )
  
      # Make sure the result has the correct shape
      self.assertEqual(len(draws), 2)
      self.assertEqual(draws[0].shape, (num_results, 3))
      self.assertEqual(draws[1].shape, (num_results, 2))
  
imports: import collections
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  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.mcmc import preconditioning_utils
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 3 for calling the DiagonalMassMatrixAdaptationShapesTest obj's underlying member method code with example usage: obj.testMultipleStateParts() without return types