jax

731e46b3 - bump jax version for pypi