code: def testMALTNamedAxis(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):
          malt_state = (
              malt.metropolis_adjusted_langevin_trajectories_init(
                  state, target_log_prob_fn=target_log_prob_fn
              )
          )
          (
              malt_state,
              malt_extra,
          ) = malt.metropolis_adjusted_langevin_trajectories_step(
              malt_state,
              damping=self._constant(0.5),
              step_size=self._constant(0.2),
              num_integrator_steps=4,
              target_log_prob_fn=target_log_prob_fn,
              named_axis=named_axis,
              seed=seed,
          )
          return malt_state, malt_extra
  
      seed = self._make_seed(_test_seed())
      malt_state, malt_extra = kernel(state, seed)
      self.assertAllClose(
          malt_state.state["shared"][0], malt_state.state["shared"][1]
      )
      self.assertTrue(
          np.any(
              np.abs(
                  malt_state.state["sharded"][0]
                  - malt_state.state["sharded"][1]
              )
              > 1e-3
          )
      )
      self.assertAllClose(
          malt_extra.is_accepted[0], malt_extra.is_accepted[1]
      )
      self.assertAllClose(
          malt_extra.log_accept_ratio[0], malt_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 malt
  from fun_mc import test_util
label_desc: Use image node_id 5 for calling the MALTTest obj's underlying member method code with example usage: obj.testMALTNamedAxis() and returns: unknown,  , malt_state,  malt_extra