jax

ce17815f - Merge pull request #34977 from mattjj:hijax-vmappable