Upgrade JAX to v0.8.2 and clean up installation instructions#1

google-labs-jules[bot] wants to merge 2 commits intojax-upgradermcao/CalCIL:jax-upgradefrom

upgrade-jax-0.8.2-5519985400734555835rmcao/CalCIL:upgrade-jax-0.8.2-5519985400734555835Copy head branch name to clipboard