jax
e702c77f - Mention keyword arguments in the vmap in_axes arity error

Commit
46 days ago
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
Author
Committer
Parents
Loading