jax

3eecb6a7 - Raise an error if an unmapped argument to vmap is sharded on mesh axis that we are vmapping over