code: def testBatches(
      self, state_part_shape, variance_part_shape, log_prob_shape
  ):
      dist = independent.Independent(
          normal.Normal(
              tf.zeros(state_part_shape), tf.ones(state_part_shape)
          ),
          reinterpreted_batch_ndims=len(state_part_shape)
          - len(log_prob_shape),
      )
      state_part = tf.zeros(state_part_shape)
  
      running_variance = sample_stats.RunningVariance.from_stats(
          num_samples=10.0,
          mean=tf.zeros(variance_part_shape),
          variance=tf.ones(variance_part_shape),
      )
  
      kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=dist.log_prob,
          num_leapfrog_steps=2,
          step_size=1.0,
      )
      kernel = dmma.DiagonalMassMatrixAdaptation(
          inner_kernel=kernel, initial_running_variance=running_variance
      )
  
      num_results = 5
      draws = sample.sample_chain(
          num_results=num_results,
          current_state=state_part,
          kernel=kernel,
          seed=test_util.test_seed(),
          trace_fn=None,
      )
  
      # Make sure the result has the correct shape
      self.assertEqual(draws.shape, (num_results,) + state_part_shape)
  
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 1 for calling the DiagonalMassMatrixAdaptationShapesTest obj's underlying member method code with example usage: obj.testBatches(state_part_shape, variance_part_shape, log_prob_shape) without return types