---
tags:
  - Free
catalog:
  name: JAX
  description: Autograd and XLA, brought together for high-performance machine learning
  license_type: Free
  disciplines:
    - Data Analytics and Machine Learning
  available_on:
    - LUMI
    - Puhti
    - Mahti
    - Roihu
---

# JAX

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

!!! info "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](https://github.com/google/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](https://lumi-supercomputer.github.io/LUMI-EasyBuild-docs/j/jax/)
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](../computing/containers/tykky.md)
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](../computing/containers/overview.md#running-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](../support/tutorials/python-usage-guide.md#installing-python-packages-to-existing-modules).
If you think that some important package should be included in the
module provided by CSC, please [contact our
servicedesk](../support/contact.md).

## License

JAX is licensed under [Apache License
2.0](https://github.com/google/jax/blob/main/LICENSE).

## Usage

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

```text
module load python-jax
```

To access JAX on Puhti or Mahti:

```text
module load jax
```

or
```text
module load jax/small
```

To access CSC-installed JAX on LUMI:

```text
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:

```bash
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](../computing/running/getting-started.md).

!!! 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](../support/tutorials/ml-data.md) for more details.

## More information

- [CSC's Machine learning guide](../support/tutorials/ml-guide.md)
- [JAX in LUMI Software Library](https://lumi-supercomputer.github.io/LUMI-EasyBuild-docs/j/jax/)
- [JAX GitHub page](https://github.com/google/jax)
- [JAX reference documentation](https://jax.readthedocs.io/en/latest/)
