NVIDIA/JAX-Toolbox

NVIDIA/JAX-Toolbox

Description: JAX-Toolbox

Language: Python

License: Apache-2.0

Stars: 412

Forks: 75

Open issues: 117

Created: 2023-02-13T14:39:12Z

Pushed: 2026-06-10T16:12:04Z

Default branch: main

Fork: no

Archived: no

README:

JAX Toolbox

 

JAX Toolbox provides a public CI, Docker images for popular JAX libraries, and optimized JAX examples to simplify and enhance your JAX development experience on NVIDIA GPUs. It supports JAX libraries such as MaxText and Pallas.

Frameworks and Supported Models

We support and test the following JAX frameworks and model architectures. More details about each model and available containers can be found in their respective READMEs.

| Framework | Models | Use cases | Container | | :--- | :---: | :---: | :---: | | maxtext| GPT, LLaMA, Gemma, Mistral, Mixtral | pre-training | ghcr.io/nvidia/jax:maxtext | | [axlearn](./docs/frameworks/axlearn/README.md) | Fuji | pre-training | ghcr.io/nvidia/jax:axlearn | | alphafold3 | EvoFormer | inference | ghcr.io/nvidia/jax:alphafold |

Build Pipeline Status

Components Container Build Test

ghcr.io/nvidia/jax:base

[no tests]

ghcr.io/nvidia/jax:jax

ghcr.io/nvidia/jax:equinox

[tests disabled]

ghcr.io/nvidia/jax:maxtext

ghcr.io/nvidia/jax:axlearn

ghcr.io/nvidia/jax:alphafold

In all cases, ghcr.io/nvidia/jax:XXX points to latest nightly build of the container for XXX. For a stable reference, use ghcr.io/nvidia/jax:XXX-YYYY-MM-DD.

In addition to the public CI, we also run internal CI nightlies on GB300, B300, GB200, B200, DGX Spark, RTX PRO 6000 Blackwell, Jetson AGX Thor, H100 SXM 80GB, A100 SXM 80GB.

Environment Variables

The JAX image is embedded with the following flags and environment variables for performance tuning of XLA and NCCL:

| XLA Flags | Value | Explanation | | --------- | ----- | ----------- | | --xla_gpu_enable_latency_hiding_scheduler | true | allows XLA to move communication collectives to increase overlap with compute kernels |

There are various other XLA flags users can set to improve performance. XLA flags can also be tuned per workload based on specific performance needs.

Versions

| First nightly with new base container | Base container | | ------------------------------------- | -------------- | | 2026-05-01 | nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 | | 2026-03-12 | nvcr.io/nvidia/cuda-dl-base:26.02-cuda13.1-devel-ubuntu24.04 | | 2025-12-17 | nvcr.io/nvidia/cuda-dl-base:25.11-cuda13.0-devel-ubuntu24.04 | | 2025-10-02 | nvcr.io/nvidia/cuda-dl-base:25.09-cuda13.0-devel-ubuntu24.04 | | 2025-08-22 | nvcr.io/nvidia/cuda-dl-base:25.08-cuda13.0-devel-ubuntu24.04 | | 2025-07-03 | nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04 | | 2025-04-11 | nvcr.io/nvidia/cuda-dl-base:25.03-cuda12.8-devel-ubuntu24.04 | | 2025-03-04 | nvcr.io/nvidia/cuda-dl-base:25.02-cuda12.8-devel-ubuntu24.04 | | 2025-01-31 | nvcr.io/nvidia/cuda-dl-base:25.01-cuda12.8-devel-ubuntu24.04 | | 2025-01-28 | nvcr.io/nvidia/cuda-dl-base:24.11-cuda12.6-devel-ubuntu24.04 | | 2024-12-07 | nvidia/cuda:12.6.3-devel-ubuntu22.04 | | 2024-11-06 | nvidia/cuda:12.6.2-devel-ubuntu22.04 | | 2024-09-25 | nvidia/cuda:12.6.1-devel-ubuntu22.04 | | 2024-07-24 | nvidia/cuda:12.5.0-devel-ubuntu22.04 |

Profiling

See [this page](./docs/profiling.md) for more information about how to profile JAX programs on GPU.

NVIDIA Staging Containers

JAX-Toolbox staging container hosts pending NVIDIA-authored XLA and JAX enhancements for NVIDIA GPU. These are pending PRs that are awaiting upstream review and merge in OSS OpenXLA and JAX repositories. For this initiative, we are publishing scale-training tagged containers - jax-scale-training.

Pipeline Container Schedule

scale-training.yaml ghcr.io/nvidia/jax:jax-scale-training Every other Saturday, 00:00 UTC

The table belowed shows all the produced scale-training containers. From Saturday 23rd May, scale-training container will be generated and published every other Saturday at 00:00 UTC. Refer to this page STAGING.md for more information on the underlying XLA staging branch, the pending PRs included. The page also lists the corresponding JAX commit used.

Staging releases: | Release date | Container | XLA branch (includes pending PRs) | | ------------------------------------- | -------------- |-------------- | | 2026-05-30 | ghcr.io/nvidia/jax:jax-scale-training-2026-05-30 | 00de6b34 | 2026-05-11 | ghcr.io/nvidia/jax:jax-scale-training-2026-05-11 | 552b0a3e | 2026-04-24 | ghcr.io/nvidia/jax:jax-scale-training-2026-04-24 | 5dfe2147 | 2026-04-18 | ghcr.io/nvidia/jax:jax-scale-training-2026-04-18 | 8147118

The underlying CUDA container, used for building -scale-training containers, can be seen from the corresponding dated version in the [Versions](#versions) tab.

To check the versions of libraries in the container, you can simply run:

docker run --rm --gpus all ghcr.io/nvidia/jax:jax-scale-training-2026-04-24 -c '
echo "=== CUDA Toolkit ===" && echo ${CUDA_VERSION}
echo "=== cuDNN ===" && echo ${CUDNN_VERSION}
echo "=== NCCL ===" && echo ${NCCL_VERSION}
echo "=== Python Packages ===" && pip list | grep -iE…