code: def test_rwm_kernel_tracks_axis_names(self):
      kernel = random_walk_metropolis.RandomWalkMetropolis(
          normal.Normal(0, 1).log_prob
      )
      self.assertIsNone(kernel.experimental_shard_axis_names)
      kernel = random_walk_metropolis.RandomWalkMetropolis(
          normal.Normal(0, 1).log_prob,
          experimental_shard_axis_names=["a"],
      )
      self.assertListEqual(kernel.experimental_shard_axis_names, ["a"])
      kernel = random_walk_metropolis.RandomWalkMetropolis(
          normal.Normal(0, 1).log_prob,
      ).experimental_with_shard_axes(["a"])
      self.assertListEqual(kernel.experimental_shard_axis_names, ["a"])
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import cauchy
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  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 random_walk_metropolis
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 1 for calling the DistributedRWMTest obj's underlying member method code with example usage: obj.test_rwm_kernel_tracks_axis_names() without return types