code: def test_vmap_both_mapped(self):
      mesh = jtu.create_global_mesh((2, 2), ("x", "y"))
      np_inp = np.arange(16).reshape(8, 2)
      s = NamedSharding(mesh, P("x", "y"))
      inp1 = jax.device_put(np_inp, s)
  
      np_inp2 = np.arange(16).reshape(2, 8)
      inp2 = jax.device_put(np_inp2, NamedSharding(mesh, P("y", "x")))
  
      @jax.jit
      def f(x, y):
          return jax.vmap(shard_alike, in_axes=(0, 1))(x, y)
  
      out1, out2 = f(inp1, inp2)
      self.assertEqual(out1.sharding, s)
      self.assertEqual(out2.sharding, s)
      self.assertArraysEqual(out1, np_inp)
      self.assertArraysEqual(out2, np_inp2.T)
  
imports: import os
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
  from jax._src import xla_bridge
  from jax._src import test_util
  from jax.sharding import NamedSharding, PartitionSpec
  from jax.experimental.shard_alike import shard_alike
  from jax.experimental.shard_map import shard_map
  from jax._src.lib import xla_extension_version
  from jax import config
label_desc: Use image node_id 13 for calling the ShardAlikeTest obj's underlying member method code with example usage: obj.test_vmap_both_mapped() and returns: jax