Repository: google/jax

Issue:
Documentation for passing in trees
I would like to pass a dictionary through a vmap. 

I made an attempt to understand https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html but couldn't figure out if this is something I need to understand or not?

`vmap(loss, [None, 0, 0])(dictionary, X[i:i+batch], y[i:i+batch])`
`vmap(loss, [None, 0, 0])(tree_flatten(dictionary), X[i:i+batch], y[i:i+batch])`
`vmap(loss, [tree_flatten(dict_none), 0, 0])(tree_flatten(dictionary), X[i:i+batch], y[i:i+batch])`

Either way I get an issue like this: 

```
ValueError: Expected list, got (([<object object at 0x7fcb21ef58b0>, <object object at 0x7fcb21ef58b0>, <object object at 0x7fcb21ef58b0>, <object object at 0x7fcb21ef58b0>, <object object at 0x7fcb21ef58b0>, <object object at 0x7fcb21ef58b0>], <object object at 0x7fcb21ef58b0>), <object object at 0x7fcb21ef58b0>, <object object at 0x7fcb21ef58b0>).

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/jax/api.py in _flatten_axes(treedef, axis_tree)
    714     msg = ("axes specification must be a tree prefix of the corresponding "
    715            "value, got specification {} for value {}.")
--> 716     raise ValueError(msg.format(axis_tree, treedef))
    717   axes = [None if a is proxy else a for a in axes]
    718   assert len(axes) == treedef.num_leaves

ValueError: axes specification must be a tree prefix of the corresponding value, got specification [([], PyTreeDef(dict[['dense1', 'dense2', 'dense3']], [PyTreeDef(dict[[]], []),PyTreeDef(dict[[]], []),PyTreeDef(dict[[]], [])])), 0, 0] for value PyTreeDef(tuple, [PyTreeDef(tuple, [PyTreeDef(list, [*,*,*,*,*,*]),*]),*,*]).
```


Discussion:
I suspect you could make a dictionary with the right keys/values by using the `fromkeys` constructor but it's a little hard to say for sure without a working example to test: `vmap(loss, [dict.fromkeys(dictionary), 0, 0])(dictionary, X[i:i+batch], y[i:i+batch])`

That said, if you don't want to include an argument in a JAX transformation, the easiest way is often just to define a new function with that argument already applied, e.g.,
```python
from functools import partial
vmap(partial(loss, dictionary))(X[i:i+batch], y[i:i+batch])
```
Thanks that's a helpful tip!
Thanks for raising this, Sasha!

We don't want to make you understand the details of pytrees; we want `vmap` and all the JAX transformations to "just work" on (nested) standard Python containers. In particular, you should _not_ need to read [the pytrees docs](https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html) (and it does say in bold at the top, "This is primarily JAX internal documentation, end-users are not supposed to need to understand this to use JAX, except when registering new user-defined container types with JAX.")

To more precisely answer your question, we might need a runnable repro.  But this works:

```python
import jax.numpy as np
from jax import vmap

dictionary = {'a': 5., 'b': np.ones(2)}
x = np.zeros(3)
y = np.arange(3.)


def f(dct, x, y):
  return dct['a'] + dct['b'] + x + y

result = vmap(f, (None, 0, 0))(dictionary, x, y)
```

Yet, understandably surprisingly, this doesn't!

```python
result = vmap(f, [None, 0, 0])(dictionary, x, y)

# ValueError: axes specification must be a tree prefix
# of the corresponding value, got specification [None, 0, 0]
# for value PyTreeDef(tuple, [PyTreeDef(dict[['a', 'b']], [*,*]),*,*]).
```

The issue is that the axis specification has to be a tree prefix of the `args` tuple, meaning an int (i.e. a kind of pytree, a leaf), a tuple (because `args` is a tuple), or a tuple of pytrees (including a tuple of ints, or a tuple of other kinds of pytrees). It can't be a list!

I think this behavior is surprising because we're so used to treating lists and tuples interchangeably in Python APIs, like we treat `'foo'` and `"foo"` string quoting interchangeably.

I want to fix this!

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -657,6 +657,14 @@ def vmap(fun: Callable, in_axes=0, out_axes=0):
   docstr = ("Vectorized version of {fun}. Takes similar arguments as {fun} "
             "but with additional array axes over which {fun} is mapped.")
 
+  if isinstance(in_axes, list):
+    # To be a tree prefix of the positional args tuple, in_axes can never be a
+    # list: if in_axes is not a leaf, it must be a tuple of trees. However,
+    # in cases like these users expect tuples and lists to be treated
+    # essentially interchangeably, so we canonicalize lists to tuples here
+    # rather than raising an error. https://github.com/google/jax/issues/2367
+    in_axes = tuple(in_axes)
+
   _check_callable(fun)
   if (not isinstance(in_axes, (list, tuple, type(None), int))
       or not isinstance(out_axes, (list, tuple, type(None), int))):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1229,6 +1229,20 @@ def vjp(x_tangent):
     b = np.dot(a + np.eye(a.shape[0]), real_x)
     print(gf(a, b))  # doesn't crash
 
+  def test_vmap_in_axes_list(self):
+    # https://github.com/google/jax/issues/2367
+    dictionary = {'a': 5., 'b': np.ones(2)}
+    x = np.zeros(3)
+    y = np.arange(3.)
+
+
+    def f(dct, x, y):
+      return dct['a'] + dct['b'] + x + y
+
+    out1 = api.vmap(f, (None, 0, 0))(dictionary, x, y)
+    out2 = api.vmap(f, [None, 0, 0])(dictionary, x, y)
+    self.assertAllClose(out1, out2, check_dtypes=True)
+
   def test_vmap_in_axes_tree_prefix_error(self):
     # https://github.com/google/jax/issues/795
     self.assertRaisesRegex(