Installing JAX is ridiculously easy, similar to NumPy or any other Python library.

CPU-only installation

If you don’t have GPU support, you can still code JAX on your CPU.

  1. Upgrade the pip: pip install --upgrade pip
  2. Install JAX: pip install --upgrade "jax[cpu]"

Get hands-on with 1400+ tech skills courses.