code: def test_kernel_can_shard_chains_across_devices(
      self,
      reduce_fn,
      skip_on_eager=False,
  ):
      if skip_on_eager and (tf.executing_eagerly() and TF_MODE):
          self.skipTest("Not supported in Eager.")
  
      def target_log_prob(a, b):
          return normal.Normal(0.0, 1.0).log_prob(
              a
          ) + sample_dist_lib.Sample(normal.Normal(a, 1.0), 4).log_prob(
              b
          )
  
      def run(seed, log_accept_ratio):
          kernel = hmc.UncalibratedHamiltonianMonteCarlo(
              target_log_prob, step_size=1e-2, num_leapfrog_steps=2
          )
          kernel = FakeMHKernel(kernel, log_accept_ratio)
          kernel = sssa.SimpleStepSizeAdaptation(
              kernel,
              10,
              reduce_fn=reduce_fn,
              experimental_reduce_chain_axis_names=self.axis_name,
          )
          sharded_kernel = sharded.Sharded(kernel, self.axis_name)
          init_seed, sample_seed = samplers.split_seed(seed)
          state_seeds = samplers.split_seed(init_seed)
          state = [
              samplers.normal(seed=state_seeds[0], shape=[]),
              samplers.normal(seed=state_seeds[1], shape=[4]),
          ]
          kr = sharded_kernel.bootstrap_results(state)
          _, kr = sharded_kernel.one_step(state, kr, seed=sample_seed)
          return kr.new_step_size
  
      seeds = self.shard_values(
          tf.stack(
              samplers.split_seed(
                  samplers.zeros_seed(), distribute_test_lib.NUM_DEVICES
              )
          ),
          0,
      )
      log_accept_ratios = self.shard_values(
          tf.convert_to_tensor([-2.0, -1.0, 0.0, 1.0])
      )
  
      step_size = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(
                  run,
                  args=(seeds, log_accept_ratios),
                  axis_name=self.axis_name,
              ),
              0,
          )
      )
  
      for i in range(distribute_test_lib.NUM_DEVICES):
          self.assertAllClose(step_size[0], step_size[i])
  
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 1 for calling the DistributedSimpleStepSizeAdaptationTest obj's underlying member method code with example usage: obj.test_kernel_can_shard_chains_across_devices(reduce_fn, skip_on_eager) and returns: unknown, kr