code: def test_momentum_distribution_has_right_shard_axis_names(self):
      if not JAX_MODE:
          self.skipTest(
              "Test in TF runs into `merge_call` error: see b/178944108"
          )
  
      def target_log_prob(a, b):
          return normal.Normal(0.0, 1.0).log_prob(
              a
          ) + distribute_lib.psum(
              normal.Normal(
                  distribute_lib.pbroadcast(a, "foo"), 1.0
              ).log_prob(b),
              "foo",
          )
  
      running_variance = [
          sample_stats.RunningVariance.from_stats(
              num_samples=10.0, mean=tf.zeros([]), variance=tf.ones([])
          )
      ] * 2
  
      kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob,
          num_leapfrog_steps=2,
          step_size=1.0,
      )
      kernel = dmma.DiagonalMassMatrixAdaptation(
          inner_kernel=kernel, initial_running_variance=running_variance
      )
      kernel = kernel.experimental_with_shard_axes([[], ["foo"]])
  
      def run(seed):
          state = [
              tf.convert_to_tensor(-10.0),
              tf.convert_to_tensor(-10.0),
          ]
          kr = kernel.bootstrap_results(state)
          state, _ = kernel.one_step(state, kr, seed=seed)
          inner_results = kr.inner_results.proposed_results
          axis_names = (
              inner_results.momentum_distribution.experimental_shard_axis_names
          )
          self.assertListEqual(list(axis_names), [[], ["foo"]])
          return state
  
      self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(
                  run,
                  args=(samplers.zeros_seed(),),
                  in_axes=None,
                  axis_name="foo",
              ),
              0,
          )
      )
  
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 2 for calling the DistributedDiagonalMMATest obj's underlying member method code with example usage: obj.test_momentum_distribution_has_right_shard_axis_names() and returns: unknown, state