Repository: google/jax

Issue:
Unclear error message phrasing in pmap (and possibly vmap and soft_pmap)
_mapped_axis_size will give the an error like the following if you attempt to inappropriately pmap a computation:  "ValueError: pmap got arg 0 of rank 0 but axis to be mapped 0. The tree of ranks is:..." It is difficult to understand what the next steps are for debugging from this from the way it is phrased.


Discussion:
I agree with @cleichner. It is very hard to understand what it means and what is to be done in case this error appears.

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1320,16 +1320,17 @@ def batched_fun(*args, **kwargs):
   return batched_fun
 
 def _mapped_axis_size(tree, vals, dims, name, *, kws=False):
-  def _get_axis_size(name: str, i:int, shape: Tuple[int, ...], axis: int):
+  def _get_axis_size(name: str, shape: Tuple[int, ...], axis: int):
     try:
       return shape[axis]
     except (IndexError, TypeError) as e:
-      ranks = tree_unflatten(tree, [np.ndim(x) for x, d in zip(vals, dims)])
-      raise ValueError(f"{name} got arg {i} of rank {len(shape)} but axis to be "
-                       f"mapped {axis}. The tree of ranks is:\n{ranks}") from e
+      min_rank = axis + 1 if axis >= 0 else -axis
+      raise ValueError(f"{name} was requested to map its argument along axis {axis}, "
+                       f"which implies that its rank should be at least {min_rank}, "
+                       f"but is only {len(shape)} (its shape is {shape})") from e
 
-  mapped_axis_sizes = {_get_axis_size(name, i, np.shape(x), d)
-                       for i, (x, d) in enumerate(zip(vals, dims))
+  mapped_axis_sizes = {_get_axis_size(name, np.shape(x), d)
+                       for x, d in zip(vals, dims)
                        if d is not None}
   try:
     size, = mapped_axis_sizes


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1968,8 +1968,10 @@ def h(a, b):
         r"\(10, \[2, 2\]\)"):
       api.vmap(h, in_axes=(0, 1))(X, [U, U])
 
-    with self.assertRaisesRegex(
-        ValueError, "vmap got arg 0 of rank 0 but axis to be mapped 0"):
+    error = (r"vmap was requested to map its argument along axis 0, which "
+             r"implies that its rank should be at least 1, but is only 0 "
+             r"\(its shape is \(\)\)")
+    with self.assertRaisesRegex(ValueError, error):
       # The mapped inputs cannot be scalars
       api.vmap(lambda x: x)(1.)
 
@@ -1978,8 +1980,10 @@ def h(a, b):
       # If the output is mapped, there must be a non-None in_axes
       api.vmap(lambda x: x, in_axes=None)(jnp.array([1., 2.]))
 
-    with self.assertRaisesRegex(
-        ValueError, "vmap got arg 0 of rank 1 but axis to be mapped 1"):
+    error = (r"vmap was requested to map its argument along axis 1, which "
+             r"implies that its rank should be at least 2, but is only 1 "
+             r"\(its shape is \(2,\)\)")
+    with self.assertRaisesRegex(ValueError, error):
       api.vmap(lambda x: x, in_axes=1)(jnp.array([1., 2.]))
 
     # Error is: TypeError: only integer scalar arrays can be converted to a scalar index
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -1059,7 +1059,11 @@ def testNegativeAxes(self):
     self.assertAllClose(jax.vmap(jnp.sum, in_axes=-1)(x),
                         jnp.sum(x, axis=(0, 1)))
 
-    with self.assertRaisesRegex(ValueError, "vmap got arg 0 of rank 3 but axis to be mapped -4"):
+
+    error = (r"vmap was requested to map its argument along axis -4, which "
+             r"implies that its rank should be at least 4, but is only 3 "
+             r"\(its shape is \(3, 4, 5\)\)")
+    with self.assertRaisesRegex(ValueError, error):
       jax.vmap(jnp.sum, in_axes=-4)(x)
 
     id = lambda y: y