vmap using in_axes doesn't handle named arguments

Passing keyword arguments to a function vectorized with jax.vmap(in_axes=..., out_axes=...) does not seem to work and results in an AssertionError.

For example:

import jax
import jax.numpy as jnp

def f(a, b, c):
    return (2*a, 3*b + c)

print(jax.vmap(f, in_axes=(0, 0, None), out_axes=0)(jnp.array([1, 2]), jnp.array([2, 4]), 0.5))  # works
# (DeviceArray([2, 4], dtype=int32), DeviceArray([ 6.5, 12.5], dtype=float32))

print(jax.vmap(f, in_axes=(0, 0, None), out_axes=0)(a=jnp.array([1, 2]), b=jnp.array([2, 4]), c=0.5))  # doesn't work

results in:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
    [... skipping hidden 1 frame]

~/miniconda3/envs/lcms/lib/python3.9/site-packages/jax/_src/tree_util.py in tree_map(f, tree, is_leaf, *rest)
    166   leaves, treedef = tree_flatten(tree, is_leaf)
--> 167   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    168   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

~/miniconda3/envs/lcms/lib/python3.9/site-packages/jax/_src/tree_util.py in <listcomp>(.0)
    166   leaves, treedef = tree_flatten(tree, is_leaf)
--> 167   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    168   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

ValueError: Tuple arity mismatch: 0 != 3; tuple: ().

During handling of the above exception, another exception occurred:

AssertionError                            Traceback (most recent call last)
/tmp/ipykernel_218392/621472192.py in <module>
      8 # (DeviceArray([2, 4], dtype=int32), DeviceArray([ 6.5, 12.5], dtype=float32))
      9 
---> 10 print(jax.vmap(f, in_axes=(0, 0, None), out_axes=0)(a=jnp.array([1, 2]), b=jnp.array([2, 4]), c=0.5))  # doesn't work

    [... skipping hidden 2 frame]

~/miniconda3/envs/lcms/lib/python3.9/site-packages/jax/api_util.py in flatten_axes(name, treedef, axis_tree, kws)
    274       # message only to be about the positional arguments
    275       treedef, leaf = treedef_children(treedef)
--> 276       assert treedef_is_leaf(leaf)
    277       axis_tree, _ = axis_tree
    278     raise ValueError(f"{name} specification must be a tree prefix of the "

AssertionError: 

Version used: jax: 0.2.18

[ricardoV94]

I was hit by this today. Can we have an explicit error while the behavior is not fixed perhaps?

[valentinmace]

Got the same problem today. The empty error message makes it very difficult to understand, by chance I had a snippet that did not contain named arguments so I could compare my code, otherwise it would have been very difficult.