Error out if wsc(x, P()) is called in a Explicit mesh context

PiperOrigin-RevId: 762021159

Author: yashk2810

jax/_src/pjit.py

@@ -2703,7 +2703,13 @@ def check_shardings_are_auto(shardings_flat):
       raise ValueError(
           'The spec of NamedSharding passed to with_sharding_constraint can'
           f' only refer to Auto axes of the mesh. Got spec={s.spec} and'
-          f' mesh={mesh}')
+          f' mesh={mesh}. You probably meant to use `reshard` API?')
+
+  cur_mesh = mesh_lib.get_abstract_mesh()
+  if cur_mesh._are_all_axes_explicit:
+    raise ValueError(
+        'with_sharding_constraint cannot be used when all axes of the mesh are'
+        ' of type `Explicit`. Please use the `reshard` API.')
 
 
 def with_sharding_constraint(x, shardings):

tests/pjit_test.py

@@ -7069,8 +7069,11 @@ def test_wsc_error(self, mesh):
         "The spec of NamedSharding passed to with_sharding_constraint"):
       jax.lax.with_sharding_constraint(np.arange(8).reshape(4, 2), s)
 
-    s = NamedSharding(mesh, P())
-    jax.lax.with_sharding_constraint(np.arange(8), s)
+    with self.assertRaisesRegex(
+        ValueError,
+        'with_sharding_constraint cannot be used when all axes of the mesh are'
+        ' of type `Explicit`'):
+      jax.lax.with_sharding_constraint(np.arange(8), NamedSharding(mesh, P()))
 
     s = NamedSharding(Mesh(mesh.devices, mesh.axis_names,
                            axis_types=(AxisType.Explicit, AxisType.Auto)),