Simple functorch: vmap should handle case where in_dims is None

[zou3519]

I was talking with @samdow this morning and it would be cool for us to update simple functorch (hopefully in a simple way) to get its implementation of vmap to handle cases where only some of the tensors are being vmapped over. This would let us experiment with what exactly the batching rule of custom_vjp should look like.