Repository: google/jax

Issue:
Improve error message w/ incorrect out_axes argument to vmap
[This colab](https://colab.research.google.com/drive/1DWziQYqL-tiZHvtSiCz8QJlQtkTSqlq8?usp=sharing) shows an incorrect use of vmap where more `out_axes` arguments are passed than the number of outputs in the mapped function. The error message does not indicate that the issue is with the `out_axes` argument, as opposed to the `in_axes` argument. It would be nice to improve this so that it's clear where the issue is.




Discussion:
@mattjj @tgale96 it seems to be pretty easy to pass a boolean flag `is_in_axes` to the function below to create a custom error message for `in_axes` and `out_axes`. However, it seems a bit naive and could make usage of this function cumbersome.
https://github.com/google/jax/blob/e808681f6c096e95be8532e4e901f5d410e0fb58/jax/api_util.py#L146
@IgorWilbert That line of thinking sounds right to me! Perhaps we can just pass in a string that gets included in the error message; we've used that pattern before and it seems too work out.

I really appreciate your thoughts and advice around here!

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -867,10 +867,11 @@ def batched_fun(*args):
     args_flat, in_tree  = tree_flatten(args)
     f = lu.wrap_init(fun)
     flat_fun, out_tree = flatten_fun_nokwargs(f, in_tree)
-    in_axes_flat = flatten_axes(in_tree, in_axes)
+    in_axes_flat = flatten_axes("vmap in_axes", in_tree, in_axes)
     _ = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap")
     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
-                              lambda: flatten_axes(out_tree(), out_axes))
+                              lambda: flatten_axes("vmap out_axes", out_tree(),
+                                                   out_axes))
     return tree_unflatten(out_tree(), out_flat)
 
   return batched_fun
@@ -1152,7 +1153,7 @@ def f_pmapped(*args, **kwargs):
       dyn_args, dyn_in_axes = args, in_axes
     args, in_tree = tree_flatten((dyn_args, kwargs))
     donated_invars = donation_vector(donate_tuple, dyn_args, kwargs)
-    in_axes_flat = flatten_axes(in_tree, (dyn_in_axes, 0))
+    in_axes_flat = flatten_axes("pmap in_axes", in_tree, (dyn_in_axes, 0))
     local_axis_size = _mapped_axis_size(in_tree, args, in_axes_flat, "pmap")
     for arg in args: _check_arg(arg)
     flat_fun, out_tree = flatten_fun(f, in_tree)
@@ -1191,7 +1192,7 @@ def soft_pmap(fun: Callable, axis_name: Optional[AxisName] = None, *,
   def f_pmapped(*args, **kwargs):
     f = lu.wrap_init(fun)
     args_flat, in_tree = tree_flatten((args, kwargs))
-    in_axes_flat = flatten_axes(in_tree, (in_axes, 0))
+    in_axes_flat = flatten_axes("soft_pmap in_axes", in_tree, (in_axes, 0))
     mapped_invars = tuple(axis is not None for axis in in_axes_flat)
     axis_size = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "soft_pmap")
     for arg in args_flat: _check_arg(arg)
diff --git a/jax/api_util.py b/jax/api_util.py
--- a/jax/api_util.py
+++ b/jax/api_util.py
@@ -143,7 +143,7 @@ def _argnums_partial(dyn_argnums, fixed_args, *dyn_args, **kwargs):
   ans = yield args, kwargs
   yield ans
 
-def flatten_axes(treedef, axis_tree):
+def flatten_axes(name, treedef, axis_tree):
   # given an axis spec tree axis_tree (a pytree with integers and Nones at the
   # leaves, i.e. the Nones are to be considered leaves) that is a tree prefix of
   # the given treedef, build a complete axis spec tree with the same structure
@@ -155,10 +155,10 @@ def flatten_axes(treedef, axis_tree):
   add_leaves = lambda i, x: axes.extend([i] * len(tree_flatten(x)[0]))
   try:
     tree_multimap(add_leaves, _replace_nones(proxy, axis_tree), dummy)
-  except ValueError as e:
-    msg = ("axes specification must be a tree prefix of the corresponding "
-           "value, got specification {} for value {}.")
-    raise ValueError(msg.format(axis_tree, treedef)) from e
+  except ValueError:
+    raise ValueError(f"{name} specification must be a tree prefix of the "
+                     f"corresponding value, got specification {axis_tree} "
+                     f"for value tree {treedef}.") from None
   axes = [None if a is proxy else a for a in axes]
   assert len(axes) == treedef.num_leaves
   return axes
diff --git a/jax/interpreters/sharded_jit.py b/jax/interpreters/sharded_jit.py
--- a/jax/interpreters/sharded_jit.py
+++ b/jax/interpreters/sharded_jit.py
@@ -267,11 +267,13 @@ def wrapped(*args, **kwargs):
       raise NotImplementedError("sharded_jit over kwargs not yet supported")
     f = lu.wrap_init(fun)
     args_flat, in_tree = tree_flatten((args, kwargs))
-    in_parts_flat = tuple(flatten_axes(in_tree.children()[0], in_parts))
+    in_parts_flat = tuple(flatten_axes("sharded_jit in_parts",
+                                       in_tree.children()[0], in_parts))
     flat_fun, out_tree = flatten_fun(f, in_tree)
     # TODO(skye): having a function-typed param in a primitive seems dicey, is
     # there a better way?
-    out_parts_thunk = lambda: tuple(flatten_axes(out_tree(), out_parts))
+    out_parts_thunk = lambda: tuple(flatten_axes("sharded_jit out_parts",
+                                                 out_tree(), out_parts))
     out = sharded_call(
         flat_fun,
         *args_flat,


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1073,8 +1073,8 @@ def test_vmap_in_axes_tree_prefix_error(self):
     # https://github.com/google/jax/issues/795
     self.assertRaisesRegex(
         ValueError,
-        "axes specification must be a tree prefix of the corresponding "
-        r"value, got specification \(0, 0\) for value "
+        "vmap in_axes specification must be a tree prefix of the corresponding "
+        r"value, got specification \(0, 0\) for value tree "
         r"PyTreeDef\(tuple, \[\*\]\).",
         lambda: api.vmap(lambda x: x, in_axes=(0, 0))(jnp.ones(3))
     )
@@ -1148,7 +1148,9 @@ def h(a, b):
 
     # Error is: TypeError: only integer scalar arrays can be converted to a scalar index
     with self.assertRaisesRegex(
-        ValueError, "axes specification must be a tree prefix of the corresponding value"):
+        ValueError,
+        "vmap out_axes specification must be a tree prefix of the "
+        "corresponding value.*"):
       api.vmap(lambda x: x, in_axes=0, out_axes=(2, 3))(jnp.array([1., 2.]))
 
     with self.assertRaisesRegex(