nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: google/jax
Pull request #21219 (2024-05-14 12:04:25)

[export] Enable model replication sharding using `jax.pmap`.
[export] Enable model replication sharding using `jax.pmap`.



Patch:
diff --git a/jax/experimental/export/_export.py b/jax/experimental/export/_export.py
index 1770c3299..1af133f8f 100644
--- a/jax/experimental/export/_export.py
+++ b/jax/experimental/export/_export.py
@@ -1122,6 +1122,8 @@ def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
     num_devices = axis_context.num_devices
   elif isinstance(axis_context, sharding_impls.SPMDAxisContext):
     num_devices = axis_context.mesh.size
+  elif isinstance(axis_context, sharding_impls.ReplicaAxisContext):
+    num_devices = axis_context.axis_env.nreps
   else:
     raise NotImplementedError(type(axis_context))
   if num_devices != exported.nr_devices:


Test patch:
diff --git a/tests/export_test.py b/tests/export_test.py
index 6773cbc56..472da235a 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -949,6 +949,25 @@ class JaxExportTest(jtu.JaxTestCase):
         "non-replicated sharding annotations"):
       export.call_exported(exp)(b)
 
+  def test_call_with_different_no_of_devices_pmap(self):
+    if len(jax.devices()) < 2:
+      self.skipTest("Need at least 2 devices")
+
+    @jax.jit
+    def f_jax(x):
+      return jnp.sum(x ** 2, axis=0)
+
+    a = jnp.arange(100, dtype=jnp.float32).reshape((1, 100))
+    res_native = f_jax(a)
+    exp = get_exported(f_jax)(a)
+    self.assertEqual(exp.nr_devices, 1)
+
+    b = jnp.arange(jax.device_count() * 100, dtype=jnp.float32).reshape(
+        (-1, 1, 100)
+    )
+    res_exported = jax.pmap(export.call_exported(exp))(b)
+    self.assertAllClose(res_native, res_exported[0])
+
   def test_call_with_different_no_of_devices_error_has_sharding_constraint(self):
     if jax.device_count() < 2:
       self.skipTest("Need at least 2 devices")
Repository: google/jax
Pull request #22074 (2024-06-24 23:03:11)

Fix `device_put` of a scalar with PositionalSharding
Fix `device_put` of a scalar with PositionalSharding

Fixes https://github.com/google/jax/issues/22073



Patch:
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index cb51e81d2..acceeef86 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -635,7 +635,7 @@ class PositionalSharding(sharding.Sharding):
     return self._remake(self._devices, new_ids)
 
   def check_compatible_aval(self, aval_shape: Shape) -> None:
-    if len(aval_shape) != len(self.shape):
+    if len(aval_shape) != len(self.shape) and not self.is_fully_replicated:
       raise ValueError(
           f"Sharding {self} is only valid for values of rank "
           f"{len(self.shape)}, but was applied to a value of rank "


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index 2666c86b4..fe7c4208d 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -959,6 +959,10 @@ class ShardingTest(jtu.JaxTestCase):
     # memory kind also appears in the repr but only for TPU.
     self.assertIn('GSPMDSharding({replicated}', repr(s2))
 
+  def test_positional_sharding_fully_replicated(self):
+    sharding = PositionalSharding(jax.devices())
+    jax.device_put(jnp.array(1), sharding.replicate())  # doesn't crash
+
   @parameterized.named_parameters(
       ("mesh_x_y",              P("x", "y"),   (4, 2), (),   False),
       ("mesh_x",                P("x"),        (4, 2), (1,), False),
@@ -989,6 +993,8 @@ class ShardingTest(jtu.JaxTestCase):
     self.assertTrue(op_shardings.are_op_shardings_equal(op1, op2))
 
   def test_positional_sharding_aval_compatible(self):
+    if jax.device_count() < 2:
+      self.skipTest('Requires >=2 devices')
     sharding = PositionalSharding(jax.devices()).reshape(1, jax.device_count())
     x = jax.random.uniform(jax.random.key(42), (256, 20, 1000))
     with self.assertRaisesRegex(
Repository: google/jax
Pull request #22098 (2024-06-26 00:06:00)

remove double dots from an error message
```python
import jax

@jax.jit
def f(qiaoqiao):
  if True and qiaoqiao:
    return 1
  return 0

f(True)
```


Before:
```
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
```

After:
```
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
```


Patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 812da70fd..b0f4c97e6 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -686,7 +686,7 @@ class Tracer(typing.Array, metaclass=StrictABCMeta):
   def _error_repr(self):
     if self.aval is None:
       return f"traced array with aval {self.aval}"
-    return f"traced array with shape {raise_to_shaped(self.aval).str_short()}."
+    return f"traced array with shape {raise_to_shaped(self.aval).str_short()}"
 
   def __array__(self, *args, **kw):
     raise TracerArrayConversionError(self)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index b7eb015d7..f95d3613d 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4765,6 +4765,14 @@ class APITest(jtu.JaxTestCase):
 
     f.trace(jnp.arange(8)).lower(lowering_platforms=('tpu',))  # doesn't crash
 
+  def test_no_double_dots_in_error_message(self):
+    @jax.jit
+    def f(x):
+      return 1 if x > 0 else 0
+
+    with self.assertRaisesRegex(TracerBoolConversionError, r"with shape bool\[\]\.[^\.]"):
+      f(0)
+
 
 class RematTest(jtu.JaxTestCase):
Repository: google/jax
Pull request #22099 (2024-06-26 00:33:56)

fix error message path since #22069 allowed Literals in jaxpr builder eqns


Patch:
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 29ace0925..d4b51731d 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -1857,10 +1857,11 @@ class JaxprStackFrame:
       produced = set(eqn.outvars) & active_vars
       if produced:
         active_vars.difference_update(produced)
-        active_vars.update(eqn.invars)
+        active_vars.update({v for v in eqn.invars if type(v) is Var})
     invar_positions = [i for i, v in enumerate(self.invars) if v in active_vars]
     constvars = active_vars & set(self.constvar_to_val)
-    const_eqns = [eqn for eqn in self.eqns if set(eqn.invars) & constvars]
+    const_eqns = [eqn for eqn in self.eqns
+                  if {v for v in eqn.invars if type(v) is Var} & constvars]
     return invar_positions, const_eqns
 
 def _const_folding_and_forwarding(


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index f95d3613d..fb1d6f4cd 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4773,6 +4773,19 @@ class APITest(jtu.JaxTestCase):
     with self.assertRaisesRegex(TracerBoolConversionError, r"with shape bool\[\]\.[^\.]"):
       f(0)
 
+  def test_inlined_literals_with_error(self):
+    @jax.jit
+    def f():
+      @partial(jax.jit, inline=True)
+      def g():
+        return jnp.sin(1.)
+      if g() > 0:
+        return 1.
+      return 0.
+
+    with self.assertRaisesRegex(TracerBoolConversionError, "Attempted boolean"):
+      f()
+
 
 class RematTest(jtu.JaxTestCase):
Repository: google/jax
Pull request #22217 (2024-07-01 18:38:48)

Run single controller device_put via efficient reshard if the device_set of input and the sharding is the same. The transfer_guard in the test fails before this CL.
Run single controller device_put via efficient reshard if the device_set of input and the sharding is the same. The transfer_guard in the test fails before this CL.



Patch:
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index 4341b761d..ec7bb81af 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -330,7 +330,7 @@ def _override_get_device_assignment(sharding, *args, **kwargs):
 def _identity_fn(x):
   return x
 
-def _mcjax_reshard(x, target_sharding):
+def _different_device_order_reshard(x, target_sharding):
   from jax._src import api, array
 
   inp_sharding = x.sharding
@@ -410,7 +410,14 @@ def _device_put_sharding_impl(x, aval, device):
     if (not s.is_fully_addressable and
         isinstance(x, array.ArrayImpl) and not x.is_fully_addressable):
       assert isinstance(s, Sharding)
-      return _mcjax_reshard(x, s)
+      return _different_device_order_reshard(x, s)
+
+    if (s.is_fully_addressable and isinstance(x, array.ArrayImpl) and
+        x.is_fully_addressable and len(s.device_set) > 1 and
+        s._internal_device_list != x.sharding._internal_device_list and  # pytype: disable=attribute-error
+        s.device_set == x.sharding.device_set):
+      assert isinstance(s, Sharding)
+      return _different_device_order_reshard(x, s)
 
     if not s.is_fully_addressable:
       if ((isinstance(x, array.ArrayImpl) and not x._committed) or


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 42b2297b5..9e1ca442a 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -39,7 +39,7 @@ from jax.errors import JAXTypeError
 from jax import lax
 from jax.lax import with_sharding_constraint
 from jax._src import prng
-from jax.sharding import PartitionSpec as P
+from jax.sharding import PartitionSpec as P, Mesh
 from jax.experimental import multihost_utils
 from jax.experimental.custom_partitioning import custom_partitioning
 from jax._src import array
@@ -4220,6 +4220,26 @@ class ArrayPjitTest(jtu.JaxTestCase):
     out2 = compiled2(jnp.arange(8))
     self.assertArraysEqual(out2, np.arange(8) * 2)
 
+  def test_device_put_efficient_reshard_single_host(self):
+    if jax.device_count() < 4:
+      self.skipTest('Requires >= 4 devices')
+
+    dev = jax.devices()
+    mesh1 = Mesh(np.array([dev[0], dev[1], dev[2], dev[3]]).reshape(2, 2),
+                 ('x', 'y'))
+    mesh2 = Mesh(np.array([dev[3], dev[2], dev[1], dev[0]]).reshape(2, 2),
+                 ('x', 'y'))
+    np_inp = np.arange(16).reshape(8, 2)
+    s1 = NamedSharding(mesh1, P('x', 'y'))
+    s2 = NamedSharding(mesh2, P('x'))
+
+    x_s1 = jax.device_put(np_inp, s1)
+
+    with jax.transfer_guard('disallow_explicit'):
+      out = jax.device_put(x_s1, s2)
+    self.assertArraysEqual(out, np_inp)
+    self.assertEqual(out.sharding, s2)
+
 
 def spec_regex(s):
   return str(s).replace(r"(", r"\(").replace(r")", r"\)")
Repository: google/jax
Pull request #22357 (2024-07-09 21:35:49)

[export] Fix poly shape check for vjp function with integer valued, polymorphic output.
[export] Fix poly shape check for vjp function with integer valued, polymorphic output.



Patch:
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
index c54a9f4b6..fd6d50c7b 100644
--- a/jax/_src/export/_export.py
+++ b/jax/_src/export/_export.py
@@ -1142,7 +1142,7 @@ def call(exported: Exported) -> Callable[..., jax.Array]:
     def fix_float0_ct(ct_res, expected_aval):
       if expected_aval.dtype != dtypes.float0:
         return ct_res
-      return ad_util.zeros_like_aval(expected_aval)
+      return ad_util.zeros_like_jaxval(ct_res)
 
     ct_res_fixed = map(fix_float0_ct,
                        ct_res_flat, exp_vjp.in_avals[len(args_flat):])


Test patch:
diff --git a/tests/export_test.py b/tests/export_test.py
index 975a81f35..b269aef28 100644
--- a/tests/export_test.py
+++ b/tests/export_test.py
@@ -462,7 +462,9 @@ class JaxExportTest(jtu.JaxTestCase):
     self.assertAllClose(jax.grad(jax.grad(jax.grad(f)))(x),
                         jax.grad(jax.grad(jax.grad(f1)))(x))
 
-  def test_grad_int(self):
+  @jtu.parameterized_filterable(
+    kwargs=[dict(poly_shape=True), dict(poly_shape=False)])
+  def test_grad_int(self, poly_shape):
     def f(xi, xf):
       return (2 * xi.T, xf.T * xf.T)
 
@@ -480,7 +482,11 @@ class JaxExportTest(jtu.JaxTestCase):
     self.assertAllClose(res, (xi_ct, xf_ct))
     (f_outi_ct2, f_outf_ct2), = f_vjp2((xi_ct, xf_ct))
 
-    exp = get_exported(jax.jit(f), vjp_order=2)(xi, xf)
+    if poly_shape:
+      args = export.symbolic_args_specs([xi, xf], shapes_specs=["2, a", "a, 4"])
+    else:
+      args = (xi, xf)
+    exp = get_exported(jax.jit(f), vjp_order=2)(*args)
     fr = exp.call
 
     res = fr(xi, xf)
Repository: google/jax
Pull request #22797 (2024-07-31 20:54:55)

Improve error when indexing with too many indices
before:
```python
>>> jnp.arange(10).sum()[0]
Traceback (most recent call last):
  ...
IndexError: Too many indices for array: 1 non-None/Ellipsis indices for dim 0.
```
after:
```python
>>> jnp.arange(10).sum()[0]
Traceback (most recent call last):
  ...
IndexError: Too many indices: 0-dimensional array indexed with 1 regular index.
```


Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 7ad74c84b..d66597958 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -8286,13 +8286,14 @@ def _is_scalar(x):
   return  np.isscalar(x) or (isinstance(x, (np.ndarray, Array))
                              and np.ndim(x) == 0)
 
-def _canonicalize_tuple_index(arr_ndim, idx, array_name='array'):
+def _canonicalize_tuple_index(arr_ndim, idx):
   """Helper to remove Ellipsis and add in the implicit trailing slice(None)."""
   num_dimensions_consumed = sum(not (e is None or e is Ellipsis or isinstance(e, bool)) for e in idx)
   if num_dimensions_consumed > arr_ndim:
+    index_or_indices = "index" if num_dimensions_consumed == 1 else "indices"
     raise IndexError(
-        f"Too many indices for {array_name}: {num_dimensions_consumed} "
-        f"non-None/Ellipsis indices for dim {arr_ndim}.")
+        f"Too many indices: {arr_ndim}-dimensional array indexed "
+        f"with {num_dimensions_consumed} regular {index_or_indices}.")
   ellipses = (i for i, elt in enumerate(idx) if elt is Ellipsis)
   ellipsis_index = next(ellipses, None)
   if ellipsis_index is not None:


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index d18244062..cbb8e92ed 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -1206,7 +1206,11 @@ class IndexingTest(jtu.JaxTestCase):
   def testWrongNumberOfIndices(self):
     with self.assertRaisesRegex(
         IndexError,
-        "Too many indices for array: 2 non-None/Ellipsis indices for dim 1."):
+        "Too many indices: 0-dimensional array indexed with 1 regular index."):
+      jnp.array(1)[0]
+    with self.assertRaisesRegex(
+        IndexError,
+        "Too many indices: 1-dimensional array indexed with 2 regular indices."):
       jnp.zeros(3)[:, 5]
Repository: google/jax
Pull request #22816 (2024-08-01 12:46:29)

Fix issue with multiple arguments when using custom_vjp with optimize_remat
This was a silly bug in how we were handling the fact that the `fwd` function expects `bool` entries for symbolic zeros. At least now I've added a test!


Patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 5c97a3dcd..90e62094f 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -1468,6 +1468,7 @@ def optimize_remat_of_custom_vjp_fwd(
     flat_fun, out_type = _flatten_fun_nokwargs(f_, in_tree)
     flat_fwd, out_trees = _flatten_fwd(fwd_, False, primal_name, fwd_name,
                                        in_tree, out_type)
+    flat_fwd = _fix_fwd_args(flat_fwd)
 
     in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat]
     fwd_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(flat_fwd, in_avals)
@@ -1497,6 +1498,12 @@ def optimize_remat_of_custom_vjp_fwd(
 
   return wrapped_fwd
 
+@lu.transformation
+def _fix_fwd_args(*args):
+  args = [(x, True) for x in args]
+  args = [x for pair in args for x in pair]
+  yield (yield args, {})
+
 def _remat_opt_impl(
     *args,
     num_consts: int,


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index f5f8075af..cb0d7c0d4 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9730,6 +9730,25 @@ class CustomVJPTest(jtu.JaxTestCase):
     v, g = jax.value_and_grad(temp)(3.2)
     self.assertAllClose(v, jnp.tan(3.2)**2)
 
+  def test_optimize_remat_multiple_args(self):
+    def f_(x, y):
+      return jnp.sin(x) * y
+
+    @jax.custom_vjp
+    def f(x, y):
+      return f_(x, y)
+
+    def f_fwd(x, y):
+      return f(x, y), (jnp.cos(x), jnp.sin(x), y)
+
+    def f_bwd(res, g):
+      cos_x, sin_x, y = res
+      return (cos_x * g * y, sin_x * g)
+
+    f.defvjp(f_fwd, f_bwd, optimize_remat=True)
+    x, y = 3.2, 1.0
+    self.assertAllClose(jax.grad(f)(x, y), jax.grad(f_)(x, y))
+
 
 def transpose_unary(f, x_example):
   def transposed(y):
Repository: google/jax
Pull request #22937 (2024-08-08 13:11:07)

Improve error message for missing vmap rule in custom_vmap
This is a partial re-land of https://github.com/google/jax/pull/22869 after it was rolled back to fix internal users. This part of the change didn't cause the issues, and I'll follow up with the rest of the changes in a second PR.


Patch:
diff --git a/jax/_src/custom_batching.py b/jax/_src/custom_batching.py
index 4d41849b7..1d405c4e5 100644
--- a/jax/_src/custom_batching.py
+++ b/jax/_src/custom_batching.py
@@ -65,6 +65,11 @@ class custom_vmap:
   @traceback_util.api_boundary
   def __call__(self, *args, **kwargs):
     assert not kwargs
+    fun_name = getattr(self.fun, "__name__", str(self.fun))
+    if not self.vmap_rule:
+      raise AttributeError(
+          f"No batching rule defined for custom_vmap function {fun_name} "
+          "using def_vmap.")
     args_flat, in_tree = tree_flatten(args)
     flat_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(self.fun), in_tree)
     in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat]


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index cb0d7c0d4..f04c7d307 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -10780,7 +10780,6 @@ class CustomVmapTest(jtu.JaxTestCase):
     )
     self.assertAllClose(outputs['b'], expected)
 
-
   def test_batch_divides_axis(self):
     def f(t):
       x, a = t
@@ -10798,6 +10797,15 @@ class CustomVmapTest(jtu.JaxTestCase):
 
     self.assertAllClose(y, (x + a)**2)
 
+  def test_undefined_rule(self):
+    @jax.custom_batching.custom_vmap
+    def f(x): return jnp.sin(x)
+
+    with self.assertRaisesRegex(
+        AttributeError, "No batching rule defined for custom_vmap function f"):
+      f(0.5)
+
+
 class CustomApiTest(jtu.JaxTestCase):
   """Test interactions among the custom_{vmap,jvp,vjp,transpose,*} APIs"""
Repository: google/jax
Pull request #22945 (2024-08-08 14:14:12)

Remove unnecessary constraint on keyword-only arguments in `custom_vjp` with `optimize_remat=True`.
Remove unnecessary constraint on keyword-only arguments in `custom_vjp` with `optimize_remat=True`.



Patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index d27b0efc7..56accc273 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -1451,7 +1451,9 @@ def optimize_remat_of_custom_vjp_fwd(
     # above and it would be good to consolidate it.
     primal_name = getattr(fun, "__name__", str(fun))
     fwd_name = getattr(fwd, "__name__", str(fwd))
-    args = _resolve_kwargs(fwd, args, kwargs)
+    # Note: we use `fun` instead of `fwd` here for consistency with
+    # custom_vjp.__call__ above.
+    args = _resolve_kwargs(fun, args, kwargs)
     if nondiff_argnums:
       for i in nondiff_argnums: _check_for_tracers(args[i])
       nondiff_argnums_ = set(nondiff_argnums)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index 15c4c8e7a..fa2f389f4 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9762,6 +9762,23 @@ class CustomVJPTest(jtu.JaxTestCase):
     x, y = 3.2, 1.0
     self.assertAllClose(jax.grad(f)(x, y), jax.grad(f_)(x, y))
 
+  def test_optimize_remat_kwargs(self):
+    @jax.custom_vjp
+    def f(x, y):
+      return jnp.sin(x) * y
+
+    def f_fwd(x, y, *, keyword=False):
+      del keyword
+      return f(x, y), (jnp.cos(x), jnp.sin(x), y)
+
+    def f_bwd(res, g):
+      cos_x, sin_x, y = res
+      return (cos_x * g * y, sin_x * g)
+
+    f.defvjp(f_fwd, f_bwd, optimize_remat=True)
+    x, y = 3.2, 1.0
+    jax.grad(f)(x, y)  # Doesn't error
+
 
 def transpose_unary(f, x_example):
   def transposed(y):
Repository: google/jax
Pull request #22973 (2024-08-09 19:02:40)

Improves error message in case of invalid sharding mesh
Improves error message in case of invalid sharding mesh



Patch:
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index 98fd8c7b0..1a23f4ba7 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -53,18 +53,17 @@ class TransferToMemoryKind:
 
 @util.cache(max_size=128, trace_context_in_key=False)
 def _check_mesh_resource_axis(mesh, parsed_pspec, _manual_axes):
-  try:
-    for p in parsed_pspec:
-      if p is not None:
-        for r in p:
-          mesh.shape[r]
-          if r in _manual_axes:
-            raise ValueError(
-                f"Axis: {r} of {parsed_pspec.get_partition_spec()} "
-                f"is also found in manual_axes: {_manual_axes}.") from None
-  except KeyError as e:
-    raise ValueError(f"Resource axis: {e.args[0]} of {parsed_pspec.user_spec} is "
-                     "undefined.") from None
+  for p in parsed_pspec:
+    if p is not None:
+      for r in p:
+        if r not in mesh.shape:
+          raise ValueError(
+              f"Resource axis: {r} of {parsed_pspec.get_partition_spec()} "
+              f"is not found in mesh: {tuple(mesh.shape.keys())}.")
+        if r in _manual_axes:
+          raise ValueError(
+              f"Axis: {r} of {parsed_pspec.get_partition_spec()} "
+              f"is also found in manual_axes: {_manual_axes}.") from None
 
 
 def hashed_index(x) -> int:


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 915a6d4b5..b93394409 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4398,7 +4398,7 @@ class PJitErrorTest(jtu.JaxTestCase):
     spec = P(resources,)
     with self.assertRaisesRegex(
         ValueError,
-        r"Resource axis: x of.*" + spec_regex(spec) + " is undefined"):
+        r"Resource axis: x of.*" + spec_regex(spec) + r" is not found in mesh: \(.*\)."):
       pjit(lambda x: x, in_shardings=spec, out_shardings=None)(x)
 
   @check_1d_2d_mesh(set_mesh=False)
@@ -4408,7 +4408,7 @@ class PJitErrorTest(jtu.JaxTestCase):
     spec = P(resources,)
     with self.assertRaisesRegex(
         ValueError,
-        r"Resource axis: x of.*" + spec_regex(spec) + " is undefined"):
+        r"Resource axis: x of.*" + spec_regex(spec) + r" is not found in mesh: \(.*\)."):
       pjit(lambda x: x, in_shardings=None, out_shardings=spec)(x)
 
   @check_1d_2d_mesh(set_mesh=False)
@@ -4418,7 +4418,7 @@ class PJitErrorTest(jtu.JaxTestCase):
     spec = P(resources,)
     with self.assertRaisesRegex(
         ValueError,
-        r"Resource axis: x of.*" + spec_regex(spec) + " is undefined"):
+        r"Resource axis: x of.*" + spec_regex(spec) + r" is not found in mesh: \(.*\)."):
       pjit(
           lambda x: with_sharding_constraint(x, spec),
           in_shardings=None,