code: def testPHMCNamedAxis(self):
      if BACKEND != "backend_jax":
          self.skipTest("JAX-only")
  
      state = {
          "sharded": jnp.zeros([4, 1024], self._dtype),
          "shared": jnp.zeros([1024], self._dtype),
      }
      in_axes = {
          "sharded": 0,
          "shared": None,
      }
      named_axis = {
          "sharded": "named_axis",
          "shared": None,
      }
  
      def target_log_prob_fn(sharded, shared):
          return (
              -(
                  backend.distribute_lib.psum(
                      jnp.square(sharded), "named_axis"
                  )
                  + jnp.square(shared)
              ),
              (),
          )
  
      @functools.partial(
          real_jax.pmap, in_axes=(in_axes, None), axis_name="named_axis"
      )
      def kernel(state, seed):
          phmc_state = prefab.persistent_hamiltonian_monte_carlo_init(
              state, target_log_prob_fn=target_log_prob_fn
          )
          (
              phmc_state,
              phmc_extra,
          ) = prefab.persistent_hamiltonian_monte_carlo_step(
              phmc_state,
              step_size=self._constant(0.2),
              num_integrator_steps=4,
              noise_fraction=0.5,
              mh_drift=0.1,
              target_log_prob_fn=target_log_prob_fn,
              named_axis=named_axis,
              seed=seed,
          )
          return phmc_state, phmc_extra
  
      seed = self._make_seed(_test_seed())
      phmc_state, phmc_extra = kernel(state, seed)
      self.assertAllClose(
          phmc_state.state["shared"][0], phmc_state.state["shared"][1]
      )
      self.assertTrue(
          np.any(
              np.abs(
                  phmc_state.state["sharded"][0]
                  - phmc_state.state["sharded"][1]
              )
              > 1e-3
          )
      )
      self.assertAllClose(
          phmc_extra.is_accepted[0], phmc_extra.is_accepted[1]
      )
      self.assertAllClose(
          phmc_extra.log_accept_ratio[0], phmc_extra.log_accept_ratio[1]
      )
  
imports: import functools
  import os
  import jax
  from jax import config
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import prefab
  from fun_mc import test_util
  from fun_mc import util_tfp
label_desc: Use image node_id 11 for calling the PrefabTest obj's underlying member method code with example usage: obj.testPHMCNamedAxis() and returns: unknown,  , phmc_state,  phmc_extra