jax

e785f894 - Raise a good error message when mesh is not provided to jax.jit when using spmd_axis_name parameter of jax.vmap