Skip to content

Docs CSC now features an automatic Finnish translation. Click here for more information.

Warning!

Puhti and Mahti are being decommissioned in stages, and their storage areas will become fully unavailable from 15 October 2026. Clean up unnecessary files and move any data you need to keep by 31 August 2026. See the Roihu data migration guide for instructions on transferring your data to Roihu.

Puhti scratch is very full: keep only active data there and move or delete everything else. No new Puhti scratch quota will be granted.

JAX

JAX is Autograd and XLA, brought together for high-performance machine learning research.

News

25.6.2026 JAX is now available on Roihu-GPU, the module has been renamed python-jax.

Available

Currently supported JAX versions:

Version Module Roihu-GPU Puhti Mahti LUMI Notes
0.10.2 python-jax/0.10 X - - - Default on Roihu-GPU
0.5.0 jax/0.5.0 default default - all packages
jax/0.5.0-small X X - framework only
0.4.38 jax/0.4.38 X X default* all packages
jax/0.4.38-small X X X* framework only
0.4.30 jax/0.4.30 X X X* all packages
jax/0.4.30-small X X X* framework only
0.4.23 jax/0.4.23-py3.9 X X X*
0.4.20 jax/0.4.20 X X X*
0.4.18 jax/0.4.18 - - X*
0.4.14 jax/0.4.14 X X -
0.4.13 jax/0.4.13 X X -
0.4.1 jax/0.4.1 X X -
0.3.13 jax/0.3.13 X X -

The modules contain JAX for Python with GPU support via CUDA/ROCm as well as a large list of additional python packages commonly used together with JAX. The Python version differs between modules releases. Modules with JAX version up to 0.4.23 use Python 3.9, later modules use Python 3.12.

Versions in LUMI, marked as "*" are still experimental with limited support. Some of the features described below might not work for them. Note that JAX is also available in the LUMI Software Library maintained by the LUMI User Support Team (and not CSC research support).

Since version 0.4.30, the JAX module comes in two flavours:

  • A "small" flavour that includes JAX with CUDA 12.2 GPU support, as well as cpu-only versions of PyTorch and Tensorflow to allow using their data loading utilities. These follow the naming scheme jax/<version>-small, or you can use jax/small to load the default (latest) version.
  • A "full" flavour that includes many commonly used machine learning packages building on JAX -- you can check pip list for a full list of all included packages. These follow the naming scheme jax/<version>, or you can simply use jax to load the default (latest) version.

On Roihu-GPU there is only a single module with a full set of packages. The module on Roihu-GPU has been renamed to python-jax.

All modules are based on containers using Apptainer (previously known as Singularity). Wrapper scripts have been provided via tykky so that common commands such as python, python3, pip and pip3 and those provided by installed packages should work as normal. For other commands, you may need to prefix them with apptainer_wrapper exec. For more information, see CSC's general instructions on how to run Apptainer containers.

Additional packages

If you find that some package is missing, you can often install it yourself using pip install. It is recommended to use Python virtual environments. See our Python documentation for more information on how to install packages yourself. If you think that some important package should be included in the module provided by CSC, please contact our servicedesk.

License

JAX is licensed under Apache License 2.0.

Usage

To use the default version of JAX on Roihu-GPU, initialize it with:

module load python-jax

To access JAX on Puhti or Mahti:

module load jax

or 

module load jax/small

To access CSC-installed JAX on LUMI:

module use /appl/local/csc/modulefiles/
module load jax  # jax/small

Note

In your scripts we recommend to fix the version so that changes in future upgrades do not break scripts, e.g.,: module load jax/0.5.0 or module load jax/0.5.0-small

Please note that the JAX modules already include the corresponding CUDA and cuDNN or ROCm libraries, so there is no need to load any cuda, cudnn, or rocm modules separately!

This will show all available versions of JAX:

module avail python-jax  # on Roihu-GPU
module avail jax         # on other systems

Note

Note that the login nodes are not intended for heavy computing, please use slurm batch jobs instead. See our instructions on how to use the batch job system.

Note

Please do not read a huge number of files from the shared file system, use fast local disk or package your data into larger files instead! See the Data storage section in our machine learning guide for more details.

More information