Introduction to JAX and Deep Learning/

...

/

Installation

Installation

Let's install JAX on the local machine and get working!

We'll cover the following...

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]"
pip install --upgrade pip
pip install --upgrade "jax[cpu]"

GPU installation

Like other deep learning systems, JAX includes support for CUDA, though we have to install the CUDA and CuDNN ourselves.

Installing CUDA

Since CUDA depends on a number of parameters − ...