Repository: jax-ml/jax
Pull request: vmap improvements
The aim here is to close #183 (a one-liner, after all this time!), close #705, and close #795.

Re: #183, we just had to replace this line:

```python
  in_tracers = map(partial(BatchTracer, trace), in_vals, in_dims)
```

with this one:

```python
  in_tracers = [BatchTracer(trace, val, dim) if dim is not None else val
                for val, dim in zip(in_vals, in_dims)]
```

Re: #705, the error printed out in that example case is now:

```python
import jax.numpy as np
from jax import api
import numpy as onp

def h(a, b):
  return np.sum(a) + np.sum(b)

X = onp.random.randn(10, 4)
U = onp.random.randn(10, 2)
api.vmap(h, in_axes=(0, 1))(X, U)
```

```
ValueError: vmap got inconsistent sizes for array axes to be mapped:
arg 0 has shape (10, 4) and axis 0 is to be mapped
arg 1 has shape (10, 2) and axis 1 is to be mapped
so
arg 0 has an axis to be mapped of size 10
arg 1 has an axis to be mapped of size 2
```

Or slightly more generally:

```python
vmap(lambda x, y, z: None, (0, 1, 0))(X, U, X)
```

```
ValueError: vmap got inconsistent sizes for array axes to be mapped:
arg 0 has shape (10, 4) and axis 0 is to be mapped
arg 1 has shape (10, 2) and axis 1 is to be mapped
arg 2 has shape (10, 4) and axis 0 is to be mapped
so
args 0, 2 have axes to be mapped of size 10
arg 1 has an axis to be mapped of size 2
```

There was a semi-interesting failure of `BatchingTest.testDynamicSlice` on eae47b2: because `vmap` now passes through un-mapped arguments unmolsted (to fix #183) that meant we couldn't write `x[i]` anymore when `x` would be a raw ndarray and `i` would be one of our Tracers. The fix was just to use `np` rather than `onp` in the test.

Problem statement:
Allow vmap to accept unbatchable arguments if they aren't mapped over
I want to use vmap on functions that have some arguments which don't need to be mapped over.  I can specify this using in_axes, but vmap doesn't allow this for some argument types.  It seems to check if every argument is a valid type for batching, even if that argument is ignored due to in_axes.  Can you remove this check for those arguments which won't be batched over?

A simple example:

```Python
from jax.api import vmap
import jax.numpy as np

def mymap(f, x):
    return f(x)

vmap(mymap, in_axes=(None, 0))(np.sin, np.linspace(0, 1, 10))
```

Fails with `TypeError: <type 'function'> is not a valid type for batching`

This is a minor issue, because I could get around with with `partial` or by making closures, but that clutters up the code a bit.

Better error messages for vmap dimensions
This piece of code is incorrect since in_axes should be (0, 0) instead of (0, 1)

def h(a, b):
 return np.sum(a) + np.sum(b)

X = onp.random.randn(10, 4)
U = onp.random.randn(10, 2)
vmap(h, in_axes=(0, 1))(X, U)

The current error message is:
"got inconsistent map dimension sizes: set([10, 2])"
 
I think the error message could describe the input shapes and mention which dimensions we are batching over.





vmap should check in_axes tree-prefix property, raise clear error
Also we can probably improve the docstring's explanation of what tree prefixes are about. Thanks,  @zhongwen !

Hints from issue discussion:
I think we can make this work by making an in-axis of `None` do something similar to what `argnums` does in `grad` (i.e. basically doing the `partial`/closure thing).
I just got this error:
```
TypeError: got inconsistent map dimension sizes: {64, 1, 10, 4}
```
which really stumped me for a while... More clarity here would definitely be helpful!

Commit message: improve vmap error messages

fixes #705

Patch to review:
diff --git a/jax/api.py b/jax/api.py
index 75c6170bf9be..853612219385 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -565,21 +565,28 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   Args:
     fun: Function to be mapped over additional axes.
-    in_axes: Specifies which input axes to map over. Normally this is a tuple with
-       one axes specification for each function argument. An integer is interpreted
-       as a tuple with the same value for all arguments. One argument axes specification
-       can be an integer (0 means first dimension), None (means that the dimension is
-       broadcasted). If the argument is a tuple of values, then the axes specification
-       can be a matching tuple as well.
-
-    out_axes: Specifies which output axes to map over. These may be integers,
-      `None`, or (possibly nested) tuples of integers or `None`.
+    in_axes: A nonnegative integer, None, or (nested) standard Python container
+      (tuple/list/dict) thereof specifying which input array axes to map over.
+      If each positional argument to ``fun`` is an array, then ``in_axes`` can
+      be a nonnegative integer, a None, or a tuple of integers and Nones with
+      length equal to the number of positional arguments to ``fun``. An integer
+      or None indicates which array axis to map over for all arguments (with
+      None indicating not to map any axis), and a tuple indicates which axis to
+      map for each corresponding positional argument. More generally, if the
+      positinal arguments to ``fun`` are container types, the corresponding
+      element of ``in_axes`` can itself be a matching container, so that
+      distinct array axes can be mapped for different container elements. The
+      constraint is that ``in_axes`` must be a container tree prefix of the
+      positional argument tuple passed to ``fun``.
+    out_axes: A nonnegative integer, None, or (nested) standard Python container
+      (tuple/list/dict) thereof indicating where the mapped axis should appear
+      in the output.
 
   Returns:
-    Batched/vectorized version of `fun` with arguments that correspond to those
-    of `fun`, but with extra array axes at positions indicated by `in_axes`, and
-    a return value that corresponds to that of `fun`, but with extra array axes
-    at positions indicated by `out_axes`.
+    Batched/vectorized version of ``fun`` with arguments that correspond to
+    those of ``fun``, but with extra array axes at positions indicated by
+    ``in_axes``, and a return value that corresponds to that of ``fun``, but
+    with extra array axes at positions indicated by ``out_axes``.
 
   For example, we can implement a matrix-matrix product using a vector dot
   product:
@@ -588,12 +595,47 @@ def vmap(fun, in_axes=0, out_axes=0):
   >>> mv = vmap(vv, (0, None), 0)      #  ([b,a], [a]) -> [b]      (b is the mapped axis)
   >>> mm = vmap(mv, (None, 1), 1)      #  ([b,a], [a,c]) -> [b,c]  (c is the mapped axis)
 
-  Here we use `[a,b]` to indicate an array with shape (a,b). Here are some
+  Here we use ``[a,b]`` to indicate an array with shape (a,b). Here are some
   variants:
 
   >>> mv1 = vmap(vv, (0, 0), 0)   #  ([b,a], [b,a]) -> [b]        (b is the mapped axis)
   >>> mv2 = vmap(vv, (0, 1), 0)   #  ([b,a], [a,b]) -> [b]        (b is the mapped axis)
   >>> mm2 = vmap(mv2, (1, 1), 0)  #  ([b,c,a], [a,c,b]) -> [c,b]  (c is the mapped axis)
+
+  Here's an example of using container types in ``in_axes`` to specify which
+  axes of the container elements to map over:
+
+  >>> A, B, C, D = 2, 3, 4, 5
+  >>> x = np.ones((A, B))
+  >>> y = np.ones((B, C))
+  >>> z = np.ones((C, D))
+  >>> def foo(tree_arg):
+  ...   x, (y, z) = tree_arg
+  ...   return np.dot(x, np.dot(y, z))
+  >>> tree = (x, (y, z))
+  >>> print(foo(tree))
+  [[12. 12. 12. 12. 12.]
+   [12. 12. 12. 12. 12.]]
+  >>> from jax import vmap
+  >>> K = 6  # batch size
+  >>> x = np.ones((K, A, B))  # batch axis in different locations
+  >>> y = np.ones((B, K, C))
+  >>> z = np.ones((C, D, K))
+  >>> tree = (x, (y, z))
+  >>> vfoo = vmap(foo, in_axes=((0, (1, 2)),))
+  >>> print(vfoo(tree))
+  [[[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]]
   """
   docstr = ("Vectorized version of {fun}. Takes similar arguments as {fun} "
             "but with additional array axes over which {fun} is mapped.")
@@ -605,12 +647,43 @@ def vmap(fun, in_axes=0, out_axes=0):
            "or a (nested) tuple of those types, got {} and {} respectively.")
     raise TypeError(msg.format(type(in_axes), type(out_axes)))
 
+  def _check_axis_sizes(tree, vals, dims):
+    try:
+      sizes, = {x.shape[d] for x, d in zip(vals, dims) if d is not None}
+    except ValueError:
+      msg = "vmap got inconsistent sizes for array axes to be mapped:\n{}"
+      # we switch the error message based on whether args is a tuple of arrays,
+      # in which case we can produce an error message based on argument indices,
+      # or if it has nested containers.
+      # TODO(mattjj,phawkins): add a way to inspect pytree kind more directly
+      if tree == tree_flatten((core.unit,) * tree.num_leaves)[1]:
+        lines1 = ["arg {} has shape {} and axis {} is to be mapped"
+                  .format(i, x.shape, d) for i, (x, d) in enumerate(zip(vals, dims))]
+        sizes = collections.defaultdict(list)
+        for i, (x, d) in enumerate(zip(vals, dims)):
+          if d is not None:
+            sizes[x.shape[d]].append(i)
+        lines2 = ["{} {} {} {} to be mapped of size {}".format(
+                   "args" if len(idxs) > 1 else "arg",
+                   ", ".join(map(str, idxs)),
+                   "have" if len(idxs) > 1 else "has",
+                   "axes" if len(idxs) > 1 else "an axis",
+                   size)
+                  for size, idxs in sizes.items()]
+        raise ValueError(msg.format("\n".join(lines1 + ["so"] + lines2)))
+      else:
+        sizes = [x.shape[d] if d is not None else None for x, d in zip(vals, dims)]
+        sizes = tree_unflatten(tree, sizes)
+        raise ValueError(msg.format("the tree of axis sizes is:\n{}".format(sizes)))
+
   @wraps(fun, docstr=docstr)
   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)
-    out_flat = batching.batch(flat_fun, args_flat, _flatten_axes(in_tree, in_axes),
+    in_axes_flat = _flatten_axes(in_tree, in_axes)
+    _check_axis_sizes(in_tree, args_flat, in_axes_flat)
+    out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
                               lambda: _flatten_axes(out_tree(), out_axes))
     return tree_unflatten(out_tree(), out_flat)
 
@@ -620,7 +693,12 @@ def _flatten_axes(treedef, axis_tree):
   dummy = tree_unflatten(treedef, [object()] * treedef.num_leaves)
   axes = []
   add_leaves = lambda i, x: axes.extend([i] * len(tree_flatten(x)[0]))
-  tree_multimap(add_leaves, _replace_nones(axis_tree), dummy)
+  try:
+    tree_multimap(add_leaves, _replace_nones(axis_tree), dummy)
+  except ValueError:
+    msg = ("axes specification must be a tree prefix of the corresponding "
+           "value, got specification {} for value {}.")
+    raise ValueError(msg.format(axis_tree, treedef))
   axes = [None if a is _none_proxy else a for a in axes]
   return axes
 
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
index d2848bfe6cab..05fa06a76484 100644
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -38,8 +38,8 @@
 
 
 def batch(fun, in_vals, in_dims, out_dim_dests):
-  out_vals, out_dims = batch_fun(fun, in_vals, in_dims)
   size, = {x.shape[d] for x, d in zip(in_vals, in_dims) if d is not not_mapped}
+  out_vals, out_dims = batch_fun(fun, in_vals, in_dims)
   return map(partial(matchaxis, size), out_dims, out_dim_dests(), out_vals)
 
 def batch_fun(fun, in_vals, in_dims):
@@ -52,7 +52,8 @@ def batch_fun(fun, in_vals, in_dims):
 @transformation_with_aux
 def batch_subtrace(master, in_dims, *in_vals):
   trace = BatchTrace(master, core.cur_sublevel())
-  in_tracers = map(partial(BatchTracer, trace), in_vals, in_dims)
+  in_tracers = [BatchTracer(trace, val, dim) if dim is not None else val
+                for val, dim in zip(in_vals, in_dims)]
   outs = yield in_tracers, {}
   out_tracers = map(trace.full_raise, outs)
   out_vals, out_dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
@@ -162,8 +163,8 @@ def get_primitive_batcher(p):
   try:
     return primitive_batchers[p]
   except KeyError:
-    raise NotImplementedError(
-        "Batching rule for '{}' not implemented".format(p))
+    msg = "Batching rule for '{}' not implemented"
+    raise NotImplementedError(msg.format(p))
 
 def defvectorized(prim):
   primitive_batchers[prim] = partial(vectorized_batcher, prim)
diff --git a/tests/api_test.py b/tests/api_test.py
index ed61d672c96d..67e0f7fe7e11 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1093,6 +1093,42 @@ 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_tree_prefix_error(self):
+    # https://github.com/google/jax/issues/795
+    jtu.check_raises_regexp(
+        lambda: api.vmap(lambda x: x, in_axes=(0, 0))(np.ones(3)),
+        ValueError, "axes specification must be a tree prefix")
+
+  def test_vmap_objects_issue_183(self):
+    # https://github.com/google/jax/issues/183
+    fun = lambda f, x: f(x)
+    vfun = api.vmap(fun, (None, 0))
+    ans = vfun(lambda x: x + 1, np.arange(3))
+    self.assertAllClose(ans, onp.arange(1, 4), check_dtypes=False)
+
+  def test_vmap_error_message_issue_705(self):
+    # https://github.com/google/jax/issues/705
+    def h(a, b):
+      return np.sum(a) + np.sum(b)
+
+    X = onp.random.randn(10, 4)
+    U = onp.random.randn(10, 2)
+    self.assertRaisesRegex(
+        ValueError,
+        "vmap got inconsistent sizes for array axes to be mapped:\n"
+        "arg 0 has shape \(10, 4\) and axis 0 is to be mapped\n"
+        "arg 1 has shape \(10, 2\) and axis 1 is to be mapped\n"
+        "so\n"
+        "arg 0 has an axis to be mapped of size 10\n"
+        "arg 1 has an axis to be mapped of size 2",
+        lambda: api.vmap(h, in_axes=(0, 1))(X, U))
+
+    self.assertRaisesRegex(
+        ValueError,
+        "vmap got inconsistent sizes for array axes to be mapped:\n"
+        "the tree of axis sizes is:\n",
+        lambda: api.vmap(h, in_axes=(0, 1))(X, [U, U]))
+
 
 if __name__ == '__main__':
   absltest.main()


Code review comments:
diff_hunk: @@ -565,21 +565,28 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   Args:
     fun: Function to be mapped over additional axes.
-    in_axes: Specifies which input axes to map over. Normally this is a tuple with
-       one axes specification for each function argument. An integer is interpreted
-       as a tuple with the same value for all arguments. One argument axes specification
-       can be an integer (0 means first dimension), None (means that the dimension is
-       broadcasted). If the argument is a tuple of values, then the axes specification
-       can be a matching tuple as well.
-
-    out_axes: Specifies which output axes to map over. These may be integers,
-      `None`, or (possibly nested) tuples of integers or `None`.
+    in_axes: A nonnegative integer, None, or (nested) standard Python container
+      (tuple/list/dict) thereof specifying which input array axes to map over.
+      If each positional argument to ``fun`` is an array, then ``in_axes`` can
+      be a nonnegative integer, a None, or a tuple of integers and Nones with
+      length equal to the number of positional arguments to ``fun``. An integer
+      or None indicates which array axis to map over for all arguments (with
+      None indicating not to map any axis), and a tuple indicates which axis to
+      map for each corresponding positional argument. More generally, if the
+      positinal arguments to ``fun`` are container types, the corresponding
path: jax/api.py
text: @user1:
s/positinal/positional

@author:
Done.

diff_hunk: @@ -565,21 +565,28 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   Args:
     fun: Function to be mapped over additional axes.
-    in_axes: Specifies which input axes to map over. Normally this is a tuple with
-       one axes specification for each function argument. An integer is interpreted
-       as a tuple with the same value for all arguments. One argument axes specification
-       can be an integer (0 means first dimension), None (means that the dimension is
-       broadcasted). If the argument is a tuple of values, then the axes specification
-       can be a matching tuple as well.
-
-    out_axes: Specifies which output axes to map over. These may be integers,
-      `None`, or (possibly nested) tuples of integers or `None`.
+    in_axes: A nonnegative integer, None, or (nested) standard Python container
+      (tuple/list/dict) thereof specifying which input array axes to map over.
+      If each positional argument to ``fun`` is an array, then ``in_axes`` can
+      be a nonnegative integer, a None, or a tuple of integers and Nones with
+      length equal to the number of positional arguments to ``fun``. An integer
+      or None indicates which array axis to map over for all arguments (with
+      None indicating not to map any axis), and a tuple indicates which axis to
+      map for each corresponding positional argument. More generally, if the
path: jax/api.py
text: @user1:
nit: remove "More generally"

@author:
Done.

diff_hunk: @@ -605,12 +647,43 @@ def vmap(fun, in_axes=0, out_axes=0):
            "or a (nested) tuple of those types, got {} and {} respectively.")
     raise TypeError(msg.format(type(in_axes), type(out_axes)))
 
+  def _check_axis_sizes(tree, vals, dims):
+    try:
+      sizes, = {x.shape[d] for x, d in zip(vals, dims) if d is not None}
path: jax/api.py
text: @user1:
Maybe make this a little more explicit and specific:
```python
mapped_axis_sizes = {x.shape[d] for x, d in zip(vals, dims) if d is not None}
try:
  size, = mapped_axis_sizes
except ValueError:
  ...
```

@author:
Good idea: this means less stuff under the `try` (because what if `x.shape[d]` raises a ValueError?) and also means it's friendlier to interactive debugging inspection. Done.

diff_hunk: @@ -565,21 +565,28 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   Args:
     fun: Function to be mapped over additional axes.
-    in_axes: Specifies which input axes to map over. Normally this is a tuple with
-       one axes specification for each function argument. An integer is interpreted
-       as a tuple with the same value for all arguments. One argument axes specification
-       can be an integer (0 means first dimension), None (means that the dimension is
-       broadcasted). If the argument is a tuple of values, then the axes specification
-       can be a matching tuple as well.
-
-    out_axes: Specifies which output axes to map over. These may be integers,
-      `None`, or (possibly nested) tuples of integers or `None`.
+    in_axes: A nonnegative integer, None, or (nested) standard Python container
+      (tuple/list/dict) thereof specifying which input array axes to map over.
+      If each positional argument to ``fun`` is an array, then ``in_axes`` can
+      be a nonnegative integer, a None, or a tuple of integers and Nones with
+      length equal to the number of positional arguments to ``fun``. An integer
+      or None indicates which array axis to map over for all arguments (with
+      None indicating not to map any axis), and a tuple indicates which axis to
+      map for each corresponding positional argument. More generally, if the
+      positinal arguments to ``fun`` are container types, the corresponding
path: jax/api.py
text: @user1:
```suggestion
      positional arguments to ``fun`` are container types, the corresponding
```

@author:
That was amazing. I haven't figured out how to do this on github _or_ in google.

diff_hunk: @@ -588,12 +595,47 @@ def vmap(fun, in_axes=0, out_axes=0):
   >>> mv = vmap(vv, (0, None), 0)      #  ([b,a], [a]) -> [b]      (b is the mapped axis)
   >>> mm = vmap(mv, (None, 1), 1)      #  ([b,a], [a,c]) -> [b,c]  (c is the mapped axis)
 
-  Here we use `[a,b]` to indicate an array with shape (a,b). Here are some
+  Here we use ``[a,b]`` to indicate an array with shape (a,b). Here are some
   variants:
 
   >>> mv1 = vmap(vv, (0, 0), 0)   #  ([b,a], [b,a]) -> [b]        (b is the mapped axis)
   >>> mv2 = vmap(vv, (0, 1), 0)   #  ([b,a], [a,b]) -> [b]        (b is the mapped axis)
   >>> mm2 = vmap(mv2, (1, 1), 0)  #  ([b,c,a], [a,c,b]) -> [c,b]  (c is the mapped axis)
+
+  Here's an example of using container types in ``in_axes`` to specify which
+  axes of the container elements to map over:
+
+  >>> A, B, C, D = 2, 3, 4, 5
+  >>> x = np.ones((A, B))
+  >>> y = np.ones((B, C))
+  >>> z = np.ones((C, D))
+  >>> def foo(tree_arg):
+  ...   x, (y, z) = tree_arg
+  ...   return np.dot(x, np.dot(y, z))
+  >>> tree = (x, (y, z))
+  >>> print(foo(tree))
+  [[12. 12. 12. 12. 12.]
+   [12. 12. 12. 12. 12.]]
+  >>> from jax import vmap
+  >>> K = 6  # batch size
+  >>> x = np.ones((K, A, B))  # batch axis in different locations
+  >>> y = np.ones((B, K, C))
+  >>> z = np.ones((C, D, K))
+  >>> tree = (x, (y, z))
+  >>> vfoo = vmap(foo, in_axes=((0, (1, 2)),))
+  >>> print(vfoo(tree))
+  [[[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]
+   [[12. 12. 12. 12. 12.]
+    [12. 12. 12. 12. 12.]]]
path: jax/api.py
text: @user1:
Nice example! Maybe print vfoo(tree).shape for brevity?

@author:
Done.

diff_hunk: @@ -565,21 +565,28 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   Args:
     fun: Function to be mapped over additional axes.
-    in_axes: Specifies which input axes to map over. Normally this is a tuple with
-       one axes specification for each function argument. An integer is interpreted
-       as a tuple with the same value for all arguments. One argument axes specification
-       can be an integer (0 means first dimension), None (means that the dimension is
-       broadcasted). If the argument is a tuple of values, then the axes specification
-       can be a matching tuple as well.
-
-    out_axes: Specifies which output axes to map over. These may be integers,
-      `None`, or (possibly nested) tuples of integers or `None`.
+    in_axes: A nonnegative integer, None, or (nested) standard Python container
+      (tuple/list/dict) thereof specifying which input array axes to map over.
+      If each positional argument to ``fun`` is an array, then ``in_axes`` can
+      be a nonnegative integer, a None, or a tuple of integers and Nones with
+      length equal to the number of positional arguments to ``fun``. An integer
+      or None indicates which array axis to map over for all arguments (with
+      None indicating not to map any axis), and a tuple indicates which axis to
+      map for each corresponding positional argument. More generally, if the
+      positinal arguments to ``fun`` are container types, the corresponding
+      element of ``in_axes`` can itself be a matching container, so that
+      distinct array axes can be mapped for different container elements. The
+      constraint is that ``in_axes`` must be a container tree prefix of the
path: jax/api.py
text: @user1:
nit: remove "The constraint is that"

@author:
Done.

diff_hunk: @@ -1093,6 +1093,42 @@ 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_tree_prefix_error(self):
+    # https://github.com/google/jax/issues/795
+    jtu.check_raises_regexp(
+        lambda: api.vmap(lambda x: x, in_axes=(0, 0))(np.ones(3)),
+        ValueError, "axes specification must be a tree prefix")
+
+  def test_vmap_objects_issue_183(self):
path: tests/api_test.py
text: @user1:
nit: rename to something more descriptive like `test_vmap_unbatched_object_passthrough` (not the greatest name but still more descriptive than an issue number?)

@author:
Done.