vmap kwarg handling

[zou3519]

Our kwarg handling seems inconsistent from JAX's:

import jax
import jax.numpy as jnp

def f(x, y):
    return x * y

x = jnp.ones([2])
jax.vmap(f)(x, y=x).shape # returns (2,)

import torch
import functorch

def f(x, y):
    return x * y

x = torch.ones(2)
functorch.vmap(f)(x, y=x).shape # returns (2, 2)

We should figure out what we want the semantics of kwarg handling to be. Right now it looks like we treat all kwargs as unbatched (and it is documented like that). There are a couple of options:

• we treat all kwargs as batched (this is kind of awkward for values that aren't Tensors)
• we treat all kwargs as unbatched (the current state)
• some mix of both? E.g., if an argument is a Tensor we treat it as batched but if it is not a Tensor (and not a pytree of Tensors) we treat it as unbatched.

[samdow]

From JAX docs
"Arguments passed as keywords are always mapped over their leading axis (i.e. axis index 0)."

To Richard's point, I'm not sure that I personally like this

[vfdev-5]

I think mix of both could be reasonable.
Speaking about the current vmap docs:

functorch/functorch/_src/vmap.py

Lines 185 to 186 in 20c6b32

	func (function): A Python function that takes one or more arguments.
	Must return one or more Tensors.

Maybe, it would worth to mention that args are Tensors and structures (pytrees) with Tensors. Otherwise, user may think that it can receive any kind of data (no matter vmap use-case) or list of scalars can be batched as well ...