Explorations in JAX¶
Welcome!¶
Getting JAX¶
If you haven’t already, make a new Python 3 virtual environment and then install jax and jaxlib from PyPI with pip
(jax-exploration) $ python -m pip install jax jaxlib
Dependencies for this tutorial¶
To get all the dependencies needed for this tutorial you can just install from the included requirements.txt in the top level binder/ directory of the source repository
(jax-exploration) $ python -m pip install -r binder/requirements.txt