code: def test_phmc_kernel_samples_correct_momenta_for_sharded_state(self):
      if not JAX_MODE:
          self.skipTest(
              "Test in TF runs into `merge_call` error: see b/178944108"
          )
  
      def target_log_prob(a, b):
          dist = normal.Normal(0.0, 1.0)
          return dist.log_prob(a) + dist.log_prob(b)
  
      kernel = phmc.PreconditionedHamiltonianMonteCarlo(
          target_log_prob, step_size=1.9, num_leapfrog_steps=2
      )
      sharded_kernel = kernel.experimental_with_shard_axes(
          [None, ["foo"]]
      )
  
      def run(seed):
          state = [0.0, 0.0]
          kr = sharded_kernel.bootstrap_results(state)
          _, kr = sharded_kernel.one_step(state, kr, seed=seed)
          return kr.proposed_results.initial_momentum
  
      momentum = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(
                  run,
                  args=(samplers.zeros_seed(),),
                  in_axes=None,
                  axis_name="foo",
              ),
              0,
          )
      )
  
      # Unsharded state momenta should all be equal
      for i in range(distribute_test_lib.NUM_DEVICES):
          self.assertAllClose(momentum[0][i], momentum[0][0])
      # Sharded state momenta should be different
      for i in range(distribute_test_lib.NUM_DEVICES):
          for j in range(distribute_test_lib.NUM_DEVICES):
              if i == j:
                  continue
              self.assertNotAllClose(momentum[1][i], momentum[1][j])
  
imports: import collections
  from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.internal import auto_composite_tensor
  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.internal import unnest
  from tensorflow_probability.python.mcmc import diagnostic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import simple_step_size_adaptation
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 2 for calling the DistributedPHMCTest obj's underlying member method code with example usage: obj.test_phmc_kernel_samples_correct_momenta_for_sharded_state() and returns: unknown, kr