jax

9855a2ee - Make eager shmap + check_vma=False + axis_index work with explicit mode.