Fix spmd_axis_name assert with explicit_mesh_axis in presence of multi-character mesh axis name

PiperOrigin-RevId: 845351158

Author: yashk2810

jax/_src/api.py

@@ -1191,9 +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 config.remove_size_one_mesh_axis_from_type.value:
+        mesh = get_abstract_mesh()
+        spmd_axis_name = tuple(i for i in spmd_axis_name if mesh.shape[i] != 1)
       if spmd_axis_name == explicit_mesh_axis:
         spmd_axis_name = None
       else:

tests/pjit_test.py

@@ -7275,11 +7275,11 @@ def f(x):
         f(arr)
 
   @parameterized.parameters(
-      (('x', 'y', 'z'), ('x', 'y')),
-      (('x', 'z'), 'x')
+      (('data', 'model', 'stage'), ('data', 'model')),
+      (('data', 'stage'), 'data')
   )
   @config.remove_size_one_mesh_axis_from_type(True)
-  @jtu.with_explicit_mesh((2, 2, 1), ('x', 'y', 'z'))
+  @jtu.with_explicit_mesh((2, 2, 1), ('data', 'model', 'stage'))
   def test_spmd_axis_name_explicit_mode_assert_remove_one_size(
       self, in_spec, out_spec, mesh):
     np_inp = np.arange(16).reshape(4, 2, 2)