Installing packages

Author

Marie-Hélène Burle

For this course, we have already installed the required packages, but this section is important for you when you will want to install packages on your machine or on the Alliance clusters.

On your machine

Before you install the stack, you need to check whether your machine is able to run JAX on GPU (all systems can run JAX on CPU).

Linux x86_64 Linux aarch64 Mac aarch64 Windows x86_64 Windows WSL2 x86_64
CPU yes yes yes yes yes
NVIDIA GPU yes yes n/a no experimental
Google TPU yes n/a n/a n/a n/a
AMD GPU yes no n/a no experimental
Apple GPU n/a no experimental n/a n/a
Intel GPU experimental n/a n/a no no

from the JAX documentation

On your machine (but not on the Alliance clusters), I recommend that you use uv to create a Python project with your chosen Python version and all the necessary packages. uv installs packages much faster than pip and it is able to resolve dependencies very well. It also manages Python versions.

If you want more information on this, I gave a webinar on uv in May 2025.

Create a Python project (let’s call it jaxai) and cd into it:

uv init --bare jaxai
cd jaxai

--bare creates a uv project without files that I am not interested in here such as a README and a main.py.

Install the packages:

uv add dm-pix imageio jax-ai-stack "jax[cuda13]" matplotlib polars scikit-image tqdm

Quick explanation of packages we are installing:

- dm-pix                ➔ for data augmentation
- imageio               ➔ to load images into NumPy arrays
- jax-ai-stack          ➔ installs JAX for the CPU (if not already installed for the GPU),
                                   Flax—the main NN library,
                                   Optax—optimizers & loss functions,
                                   Orbax—for checkpointing,
                                   Grain—to build efficient dataloaders,
                                   ml_dtypes—NumPy dtype extensions for deep learning
- jax[cuda13]           ➔ only if you want to run JAX on GPUs (use an appropriate version of Cuda)
- matplotlib            ➔ to display samples
- polars                ➔ for DataFrames
- scikit-image          ➔ image resizing
- tqdm                  ➔ display progress bars

If you install packages which depend on JAX (e.g. Flax), they will by default install the CPU version of JAX. If you want to run JAX on GPUs, make sure to install jax[cuda13] (or whatever Cuda version is suitable).

You will see that the dependencies have automatically populated a pyproject.toml file and that a virtual environment called .venv was created.

As long as you are within the project, you don’t need to activate that virtual environment. You can just launch Python (or IPython, ptpython, Jupyter…) and the packages will be available.

Alternatively, if your required a tighter pip equivalent, you can activate it as you would any other Python virtual environment:

source .venv/bin/activate

On an Alliance cluster

I already installed all the necessary packages in the training cluster to save time and space. The instructions for today thus differ from what you would normally do and production cluster instructions in the second tab are for your future reference only.

Look for available Python modules:

module spider python

Load version 3.11.5:

module load python/3.11.5

I created a virtual Python environment with all necessary packages under /project. All you have to do today is activate it with:

source /project/def-sponsor00/nabirds_venv/bin/activate

Look for available Python modules:

module spider python

Load the version of your choice:

module load python/3.13.2

Create a Python virtual environment:

python -m venv ~/env

Activate it:

source ~/env/bin/activate

Update pip from wheel:

python -m pip install --upgrade --no-index pip

Whenever a Python wheel for a package is available on the Alliance clusters, you should use it instead of downloading the package from PyPI. To do this, simply add the --no-index flag to the install command.

You can see whether a wheel is available with avail_wheels <package> or look at the list of available wheels.

Advantages of wheels:

  • compiled for the clusters hardware,
  • ensures no missing or conflicting dependencies,
  • much faster installation.

Install libraries from wheel:

python -m pip install --no-index dm-pix imageio jax-ai-stack "jax[cuda]" matplotlib polars scikit-image tqdm
  • Don’t forget --no-index to install from wheels.

We don’t have cuda13 on the Alliance clusters yet. If you follow JAX official installation instructions and try to install jax[cuda13], pip will default to installing jax (CPU only).

Make sure to use jax[cuda] (or the equivalent jax[cuda12]) instead to have the GPU version installed.