Improve the empty mesh error message raised in pjit if mesh is not us…

…ed and Pspec is passed to in|out_shardings

PiperOrigin-RevId: 517495400

jax/_src/pjit.py

@@ -706,10 +706,12 @@ def _create_sharding_for_array(mesh, x):
         "PartitionSpecs to in_shardings or out_shardings. Please pass in "
         "the `Sharding` explicitly via in_shardings or out_shardings.")
   if mesh.empty:
-    raise RuntimeError("pjit requires a non-empty mesh! Is a mesh defined at "
-                       "the call site? Alternatively, provide a "
-                       "XLACompatibleSharding to pjit and then the "
-                       "mesh context manager is not required.")
+    raise RuntimeError(
+        'pjit requires a non-empty mesh if you are passing `PartitionSpec`s or'
+        ' `None` to in_shardings or out_shardings! Is a mesh defined at the'
+        ' call site? Alternatively, provide `XLACompatibleSharding`s to'
+        ' `in_shardings` and `out_shardings` and then the mesh context manager'
+        ' is not required.')
   # A nice user error is raised in _prepare_axis_resources.
   assert isinstance(x, ParsedPartitionSpec), x
   return _create_mesh_pspec_sharding_from_parsed_pspec(mesh, x)

tests/pjit_test.py

@@ -1777,9 +1777,8 @@ def test_array_enabled_non_empty_mesh_with_pspec(self):
     arr = jnp.array([1, 2, 3])
     with self.assertRaisesRegex(
         RuntimeError,
-        "pjit requires a non-empty mesh!.*Alternatively, provide a "
-        "XLACompatibleSharding to pjit and then the mesh context manager is "
-        "not required."):
+        r'pjit requires a non-empty mesh if you are passing `PartitionSpec`s or'
+        r' `None` to in_shardings or out_shardings.*'):
       pjit(lambda x: x, in_shardings=P('x'))(arr)
 
     with self.assertRaisesRegex(
@@ -3080,9 +3079,9 @@ def testCatchesInnerXMapErrors(self):
       f(x, x)
 
   def testEmptyMesh(self):
-    error = (r"pjit requires a non-empty mesh!.*Alternatively, provide a "
-              "XLACompatibleSharding to "
-              r"pjit and then the mesh context manager is not required.")
+    error = (
+        r'pjit requires a non-empty mesh if you are passing `PartitionSpec`s or'
+        r' `None` to in_shardings or out_shardings.*')
     with self.assertRaisesRegex(RuntimeError, error):
       pjit(lambda x: x, in_shardings=None, out_shardings=None)(jnp.arange(4))