code: def testShouldPropagateShardAxisNames(self):
      test_kernel = sssa.SimpleStepSizeAdaptation(
          FakeMHKernel(
              FakeSteppedKernel(step_size=1.0), log_accept_ratio=0.0
          ),
          num_adaptation_steps=1,
      )
      self.assertIsNone(test_kernel.experimental_shard_axis_names)
      sharded_test_kernel = test_kernel.experimental_with_shard_axes(
          ["foo"]
      )
      self.assertListEqual(
          sharded_test_kernel.experimental_shard_axis_names, ["foo"]
      )
      sharded_inner_kernel = sharded_test_kernel.inner_kernel
      self.assertListEqual(
          sharded_inner_kernel.experimental_shard_axis_names, ["foo"]
      )
      self.assertListEqual(
          sharded_inner_kernel.inner_kernel.experimental_shard_axis_names,
          ["foo"],
      )
  
imports: import collections
  import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.mcmc import sharded
  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.math import generic
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
label_desc: Use image node_id 2 for calling the SimpleStepSizeAdaptationStaticBroadcastingTest obj's underlying member method code with example usage: obj.testShouldPropagateShardAxisNames() without return types