A flexible Python library for machine learning research See the installation directions for JAX on the Intro to the Machine Learning Libraries repo on GitHub.