code: def test_dmma_kernel_tracks_axis_names(self):
      def _make_kernel(**kwargs):
          running_variance = sample_stats.RunningVariance.from_stats(
              num_samples=10.0, mean=tf.zeros(5), variance=tf.ones(5)
          )
  
          kernel = phmc.PreconditionedHamiltonianMonteCarlo(
              target_log_prob_fn=sample_dist_lib.Sample(
                  normal.Normal(0.0, 1.0), 5
              ).log_prob,
              num_leapfrog_steps=2,
              step_size=1.0,
          )
          kernel = dmma.DiagonalMassMatrixAdaptation(
              inner_kernel=kernel,
              initial_running_variance=running_variance,
              **kwargs,
          )
          return kernel
  
      kernel = _make_kernel()
      self.assertIsNone(kernel.experimental_shard_axis_names)
      kernel = _make_kernel(experimental_shard_axis_names=["a"])
      self.assertListEqual(kernel.experimental_shard_axis_names, ["a"])
      kernel = _make_kernel().experimental_with_shard_axes(["a"])
      self.assertListEqual(kernel.experimental_shard_axis_names, ["a"])
  
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 DistributedDiagonalMMATest obj's underlying member method code with example usage: obj.test_dmma_kernel_tracks_axis_names() and returns: kernel