Remove one sized mesh axis from spmd_axis_name during comparison with explicit axes if remove_size_one_mesh_axis_from_type is turned on.

PiperOrigin-RevId: 842467979

Author: yashk2810

jax/_src/api.py

@@ -68,7 +68,7 @@
 from jax._src.lib import xla_client as xc
 from jax._src.lib import pmap_lib
 from jax._src.sharding import Sharding
-from jax._src.mesh import get_concrete_mesh
+from jax._src.mesh import get_concrete_mesh, get_abstract_mesh
 from jax._src.sharding_impls import (PmapSharding, PartitionSpec as P,
                                      NamedSharding)
 from jax._src.layout import Format
@@ -1191,6 +1191,9 @@ def vmap_f(*args, **kwargs):
                   _mapped_axis_size(fun, in_tree, args_flat, in_axes_flat, "vmap"))
     explicit_mesh_axis = _mapped_axis_spec(args_flat, in_axes_flat)
     if spmd_axis_name is not None and explicit_mesh_axis is not None:
+      spmd_axis_name = (
+          tuple(core.remove_size_one_mesh_axis(P(spmd_axis_name), get_abstract_mesh()))
+          if config.remove_size_one_mesh_axis_from_type.value else spmd_axis_name)
       if spmd_axis_name == explicit_mesh_axis:
         spmd_axis_name = None
       else:

tests/pjit_test.py

@@ -7268,6 +7268,25 @@ def f(x):
         "Only one of spmd_axis_name or arrays sharded on.*spmd_axis_name"):
         f(arr)
 
+  @config.remove_size_one_mesh_axis_from_type(True)
+  @jtu.with_explicit_mesh((2, 1), ('x', 'y'))
+  def test_spmd_axis_name_explicit_mode_assert_remove_one_size(self, mesh):
+    np_inp = np.arange(16).reshape(8, 2)
+    arr = jax.device_put(np_inp, NamedSharding(mesh, P(('x', 'y'), None)))
+
+    @jax.jit
+    @partial(jax.vmap, spmd_axis_name=('x', 'y'))
+    def f(x):
+      # breakpoint()
+      self.assertEqual(x.aval.sharding.spec, P(None))
+      out = x * 2
+      self.assertEqual(out.aval.sharding.spec, P(None))
+      return out
+
+    out = f(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
+    self.assertArraysEqual(out, np_inp * 2)
+
   @jtu.with_explicit_mesh((2,), ('x',))
   def test_unmapped_last_vmap(self, mesh):
     np_inp = np.arange(8)