Mention keyword arguments in the vmap in_axes arity error
A tuple in_axes only corresponds to positional arguments, and keyword
arguments are always mapped along axis 0 (documented behavior). Calls
that pass arguments by keyword hit the in_axes/args length check with a
message that only reported len(in_axes) and len(args), never saying that
the keyword arguments are the cause. Append a note naming the keyword
arguments, explaining the axis-0 rule, and suggesting fixes: shorten
in_axes if axis-0 mapping is intended, pass the arguments positionally,
or bind unmapped keyword arguments with functools.partial.
Message-only change; the error without keyword arguments is unchanged.
Helps with jax-ml/jax#7465