code: def test_gbtla_kernel_tracks_axis_names(self):
      inner_kernel = hmc.HamiltonianMonteCarlo(
          normal.Normal(0, 1).log_prob,
          step_size=1.9,
          num_leapfrog_steps=2,
      )
      kernel = gbtla.GradientBasedTrajectoryLengthAdaptation(
          inner_kernel, 1
      )
      self.assertIsNone(kernel.experimental_shard_axis_names)
      kernel = gbtla.GradientBasedTrajectoryLengthAdaptation(
          inner_kernel, 1, experimental_shard_axis_names=["a"]
      )
      self.assertListEqual(kernel.experimental_shard_axis_names, ["a"])
      kernel = gbtla.GradientBasedTrajectoryLengthAdaptation(
          inner_kernel, 1
      ).experimental_with_shard_axes(["a"])
      self.assertListEqual(kernel.experimental_shard_axis_names, ["a"])
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.mcmc import gradient_based_trajectory_length_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  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 dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
label_desc: Use image node_id 1 for calling the DistributedGBTLATest obj's underlying member method code with example usage: obj.test_gbtla_kernel_tracks_axis_names() without return types