Show the correct sharding in is_compatible_aval error in MeshPspecSharding when created via _from_parsed_pspec. Preserve the original PartitionSpec from ParsedPartitionSpec if it exists, else calculate it.

PiperOrigin-RevId: 473267905

Author: yashk2810

jax/experimental/pjit.py

@@ -637,6 +637,15 @@ def __init__(self, user_spec, partitions, sync=SpecSync.IN_SYNC):
   def user_spec(self):
     return self.unsynced_user_spec(SpecSync.IN_SYNC)
 
+  def get_partition_spec(self) -> PartitionSpec:
+    if self.sync < SpecSync.IN_SYNC:
+      return _get_single_pspec(self)
+    else:
+      if isinstance(self.unsafe_user_spec, PartitionSpec):
+        return self.unsafe_user_spec
+      else:
+        return _get_single_pspec(self)
+
   def unsynced_user_spec(self, min_sync):
     if self.sync < min_sync:
       raise AssertionError(f"Please open a bug report! ({self.sync} >= {min_sync})")
@@ -1283,7 +1292,7 @@ def with_sharding_constraint(x, axis_resources):
         for s in sharding_flat
     ]
   else:
-    sharding_flat = [MeshPspecSharding._from_parsed_pspec(mesh, a)
+    sharding_flat = [pxla._create_mesh_pspec_sharding(mesh, a.user_spec, a)
                      for a in axis_resources_flat]
     # Calculate unconstrained_dims from MeshPspecSharding because that information
     # is lost when converted to OpSharding. Bind unconstrained_dims to

jax/experimental/sharding.py

@@ -197,8 +197,7 @@ def is_compatible_aval(self, aval_shape: Shape):
 
   @classmethod
   def _from_parsed_pspec(cls, mesh, parsed_pspec):
-    from jax.experimental import pjit
-    return cls(mesh, pjit._get_single_pspec(parsed_pspec), parsed_pspec)
+    return cls(mesh, parsed_pspec.get_partition_spec(), parsed_pspec)
 
   @pxla.maybe_cached_property
   def device_set(self) -> Set[Device]:

tests/array_test.py

@@ -505,6 +505,20 @@ def test_pmap_sharding_hash_eq(self):
     self.assertGreater(cache_info2.hits, cache_info1.hits + 1)
     self.assertEqual(cache_info2.misses, cache_info1.misses)
 
+  def test_is_compatible_error(self):
+    shape = (8, 2)
+    mesh = jtu.create_global_mesh((1, 1, 2), ('replica', 'data', 'mdl'))
+    mps = sharding.MeshPspecSharding(mesh, P(None, ('mdl',), None, None))
+    new_mps = sharding.MeshPspecSharding._from_parsed_pspec(
+        mps.mesh, mps._parsed_pspec)
+
+    with self.assertRaisesRegex(
+        ValueError,
+        r"Sharding MeshPspecSharding\(mesh={'replica': 1, 'data': 1, 'mdl': 2}, "
+        r"partition_spec=PartitionSpec\(None, \('mdl',\), None, None\)\) is only "
+        "valid for values of rank at least 4, but was applied to a value of rank 2"):
+      new_mps.is_compatible_aval(shape)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())

tests/pjit_test.py

@@ -1052,6 +1052,24 @@ def make_keys(seeds):
       self.assertEqual(out.shape, input_shape)
       out.unsafe_raw_array()  # doesn't crash
 
+  def test_with_sharding_constraint_is_compatible_error(self):
+    mesh = jtu.create_global_mesh((1, 1, 2), ('replica', 'data', 'mdl'))
+
+    with mesh:
+      def f(x):
+        y = with_sharding_constraint(x, P(None, ('mdl',), None, None))
+        z = y + 2
+        return z
+      pjit_f = pjit(f, in_axis_resources=P(None), out_axis_resources=P(None))
+
+      with self.assertRaisesRegex(
+          ValueError,
+          r"One of with_sharding_constraint.*Sharding "
+          r"MeshPspecSharding\(mesh={'replica': 1, 'data': 1, 'mdl': 2}, "
+          r"partition_spec=PartitionSpec\(None, \('mdl',\), None, None\)\) is only "
+          "valid for values of rank at least 4, but was applied to a value of rank 1"):
+        pjit_f(jnp.array([1, 2, 3]))
+
 
 class GDAPjitTest(jtu.JaxTestCase):
 
@@ -2180,8 +2198,7 @@ def testRankTooLowConstraint(self):
     x = jnp.arange(2)
     spec = P('x', 'y')
     error = re.compile(
-        r"One of with_sharding_constraint arguments" + r".*" + spec_regex(
-            pxla.array_mapping_to_axis_resources(pxla._get_array_mapping(spec))) +
+        r"One of with_sharding_constraint arguments" + r".*" + spec_regex(spec) +
         r".*rank at least 2, but was applied to a value of rank 1", re.M | re.S)
     with self.assertRaisesRegex(ValueError, error):
       pjit(lambda x: with_sharding_constraint(x, spec),
@@ -2625,6 +2642,22 @@ def test_simulated_training_cache_in_pjit(self):
     self.assertEqual(id(next_op_sharding_sharding._op_sharding),
                      id(op_sharding_sharding._op_sharding))
 
+  def test_get_partition_spec(self):
+    mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
+    s = MeshPspecSharding(mesh, P('x', 'y', None))
+
+    self.assertEqual(s._parsed_pspec.get_partition_spec(), P('x', 'y', None))
+
+    recovered_parsed_pspec = pjit_lib.parse_flatten_op_sharding(
+        s._to_xla_op_sharding(3), mesh)
+    self.assertEqual(recovered_parsed_pspec[0].get_partition_spec(),
+                     P(('x',), ('y',)))
+
+    out_of_sync_parsed_pspec = pjit_lib.ParsedPartitionSpec(
+        P('x', 'y'), ('x', 'y'), pjit_lib.SpecSync.OUT_OF_SYNC)
+    self.assertEqual(out_of_sync_parsed_pspec.get_partition_spec(),
+                     P(('x',), ('y',)))
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())