Update jax version in requirements.#4

copybara-service[bot] wants to merge 1 commit intomaingoogle-deepmind/distrax:mainfrom

test_368589236google-deepmind/distrax:test_368589236Copy head branch name to clipboard