code: def test_momentum_dists(self):
      state_parts = [
          tf.ones([13, 5, 3]),
          tf.ones([13, 5]),
          tf.ones([13, 5, 2, 4]),
      ]
      batch_shape = [13, 5]
      md = pu.make_momentum_distribution(state_parts, batch_shape)
      md = pu.update_momentum_distribution(
          md,
          tf.nest.map_structure(
              lambda s: tf.reduce_sum(s, (0, 1)), state_parts
          ),
      )
      self.evaluate(tf.nest.flatten(md, expand_composites=True))
  
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 9 for calling the DiagonalMassMatrixAdaptationTest obj's underlying member method code with example usage: obj.test_momentum_dists() without return types