See the installation directions for JAX on the Intro to the Machine Learning Libraries repo on GitHub.