Installing packages

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

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

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

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 you need to for advanced workflow involving other tools (e.g. Quarto), 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.

Today

Look for available Python modules:

module spider python

Load the version of your choice:

TensorFlow and all packages depending from it (including TensorFlow Datasets and Grain) are still not (as of April 2025) ported to Python 3.13.

module load python/3.13.2

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

source /project/60055/env/bin/activate

Production cluster

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 pip --no-index

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[cuda13] matplotlib polars scikit-image

Don’t forget --no-index to install from wheels.