Repository: google/jax
Issue: Confusing `vmap` user experience when passing in `False` into `in_axes`

`False` is interpreted as 0 in the `in_axes` argument to `jax.vmap`. This is an invalid value (and some users may expect it to behave like `None`). See this code:

```python
import jax.numpy as jnp

def f(x, y):
  return jnp.sum(x) + jnp.sum(y)

print(jax.vmap(f, in_axes=[False, 0])(arr, arr))
print(jax.vmap(f, in_axes=[None, 0])(arr, arr))
print(jax.vmap(f, in_axes=[0, 0])(arr, arr))
```

prints

```
[ 6 14]
[13 17]
[ 6 14]
```

The `vmap` documentation states that each item in `in_axes` must be either an integer or None. @mbz was confused by this and accidentally thought that `False` would behave like `None`. But instead it's interpreted as 0.

It would be good to raise an error for users accidentally using `False` in `in_axes`.

Changed files:
jax/api.py, tests/api_test.py

diff --git a/jax/api.py b/jax/api.py
index ded33746a..0fc1fb0a4 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -1259,14 +1259,12 @@ def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F:
     # rather than raising an error. https://github.com/google/jax/issues/2367
     in_axes = tuple(in_axes)
 
-  in_axes_, out_axes_ = tree_leaves(in_axes), tree_leaves(out_axes)
-  if not all(isinstance(l, (type(None), int)) for l in in_axes_):
+  if not all(type(l) is int for l in tree_leaves(in_axes)):
     raise TypeError("vmap in_axes must be an int, None, or (nested) container "
                     f"with those types as leaves, but got {in_axes}.")
-  if not all(isinstance(l, (type(None), int)) for l in out_axes_):
+  if not all(type(l) is int for l in tree_leaves(out_axes)):
     raise TypeError("vmap out_axes must be an int, None, or (nested) container "
                     f"with those types as leaves, but got {out_axes}.")
-  del in_axes_, out_axes_
 
   @wraps(fun, docstr=docstr)
   @api_boundary
@@ -1560,6 +1558,13 @@ def pmap(
   donate_tuple = rebase_donate_argnums(_ensure_index_tuple(donate_argnums),
                                        static_broadcasted_tuple)
 
+  if not all(type(l) is int for l in tree_leaves(in_axes)):
+    raise TypeError("pmap in_axes must be an int, None, or (nested) container "
+                    f"with those types as leaves, but got {in_axes}.")
+  if not all(type(l) is int for l in tree_leaves(out_axes)):
+    raise TypeError("pmap out_axes must be an int, None, or (nested) container "
+                    f"with those types as leaves, but got {out_axes}.")
+
   @wraps(fun)
   @api_boundary
   def f_pmapped(*args, **kwargs):
diff --git a/tests/api_test.py b/tests/api_test.py
index 395ffd58f..d2f340035 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1939,6 +1939,16 @@ class APITest(jtu.JaxTestCase):
         foo, in_axes=((0, collections.OrderedDict([('a', 1), ('b', 2)])),))
     self.assertEqual(vfoo(tree).shape, (6, 2, 5))
 
+  def test_vmap_in_axes_bool_error(self):
+    # https://github.com/google/jax/issues/6372
+    with self.assertRaisesRegex(TypeError, "must be an int"):
+      api.vmap(lambda x: x, in_axes=False)(jnp.zeros(3))
+
+  def test_pmap_in_axes_bool_error(self):
+    # https://github.com/google/jax/issues/6372
+    with self.assertRaisesRegex(TypeError, "must be an int"):
+      api.pmap(lambda x: x, in_axes=False)(jnp.zeros(1))
+
   def test_pmap_global_cache(self):
     def f(x, y):
       return x, y