jax

ee312afe - Corresponding `build.py` updates after https://github.com/jax-ml/jax/commit/5132a188f75022a2ef1387d944f0285557876e24