code: def test_gbtla_kernel_computes_same_criterion_info_with_sharded_state(
      self,
      criterion_fn,
  ):
      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",
          )
  
      kernel = hmc.HamiltonianMonteCarlo(
          target_log_prob, step_size=1e-2, num_leapfrog_steps=2
      )
      kernel = gbtla.GradientBasedTrajectoryLengthAdaptation(
          kernel, 10, criterion_fn=criterion_fn
      )
      sharded_kernel = kernel.experimental_with_shard_axes(
          [None, ["foo"]]
      )
  
      def run(seed):
          init_seed, sample_seed = samplers.split_seed(seed)
          state_seeds = samplers.split_seed(init_seed)
          state = [
              samplers.normal(seed=state_seeds[0], shape=[5]),
              samplers.normal(seed=state_seeds[1], shape=[5]),
          ]
          kr = sharded_kernel.bootstrap_results(state)
          _, kr = sharded_kernel.one_step(state, kr, seed=sample_seed)
          return (
              kr.criterion,
              kr.averaged_sq_grad,
              kr.averaged_max_trajectory_length,
          )
  
      criterion, avg_sq_grad, avg_max_tl = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(
                  run,
                  args=(samplers.zeros_seed(),),
                  in_axes=None,
                  axis_name="foo",
              ),
              0,
          )
      )
  
      for i in range(distribute_test_lib.NUM_DEVICES):
          self.assertAllClose(criterion[0], criterion[i])
          self.assertAllClose(avg_sq_grad[0], avg_sq_grad[i])
          self.assertAllClose(avg_max_tl[0], avg_max_tl[i])
  
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 2 for calling the DistributedGBTLATest obj's underlying member method code with example usage: obj.test_gbtla_kernel_computes_same_criterion_info_with_sharded_state(criterion_fn) and returns: unknown, kr,  kr,  kr