JAX Toolbox
JAX Toolbox is NVIDIA's home for JAX and XLA on GPUs, including the latest updates, containers and educational resources.
Latest news | Tech blogs | Documentation | Resources | Container images
What's included
-
Latest news: New NGC JAX releases, new hardware support, major JAX and XLA optimizations and updates on the JAX on GPU stack.
-
Tech blogs: Blogs highlighting key technical challenges, practical guidance and observed results
-
Documentation: Overview of JAX on GPUs, customizing JAX with NVIDIA CUDA Toolkit and going beyond what’s possible in native JAX, practical guides for profiling, optimizing, scaling and third-party tool integrations in JAX.
-
Resources: Talks, presentations and ecosystem references
-
Container images: Nightly builds, monthly NGC releases, and staging containers for JAX on GPU projects
-
Experimental projects: Experimental projects for developers to evaluate and provide feedback on
Latest news
- NGC 26.07 JAX and MaxText images released
CUDA 13.3.1, JAXv0.10.2- Improved NCCL plugin expereince on GCP; cuDNN-backed flex attention re-enabled; support for CuDNN AoT compilation in XLA; LLM training speedup
- NGC 26.06 JAX and MaxText images released
CUDA 13.3.0, JAXv0.10.1- Improved GPU dot fusion cost model and introduced tileable GemmFusion.
- NGC 26.05 JAX and MaxText images released
CUDA 13.2.1, JAXv0.10.0- performant
jax.ragged_dot; XLA’snvfp4kernel codegen improvements; multi-stream collective support in XLA; PDL (programmatic dependent launch) enablement in XLA and improved D2H & H2D copy overlap with compute
- NGC 26.04 JAX and MaxText images released
CUDA 13.2.1, JAXv0.9.2- Thor support in JAX/XLA with full Triton; FP8 reduction support in XLA; and improved collective cost model of GB200/GB300
Detailed release notes for NGC JAX containers are published in the NGC JAX release notes.
Tech blogs
- [Jul 2026] Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading
- [Jun 2026] Train Models Faster with JAX and MaxText Using NVFP4 on NVIDIA Blackwell
- [Feb 2026] Accelerating Long-Context Model Training in JAX and XLA
- [Jul 2025] Optimizing for Low-Latency Communication in Inference Workloads with JAX and XLA
Resources
Talks and presentations
2026
- [GTC 2026] JAX on Blackwell customer success story
- [GTC 2026] Large-Scale JAX Training With Pipeline Parallelism
- [NVIDIA Developer] Two Ways to Fine-Tune JAX on NVIDIA GPUs: PEFT and SFT with Tunix and MaxText
- [NVIDIA Developer] CuTe DSL for JAX Developers: Writing Custom GPU Kernels in Python
2025
- [JAX/OpenXLA DevLab Fall 2025] Accelerating LLM models on Blackwell
- [JAX/OpenXLA DevLab Fall 2025] Leveraging CuTeDSL in JAX
- [JAX/OpenXLA DevLab Fall 2025] JaxPP: A library for MPMD training in JAX
- [GTC 2025] Horizontal Scaling of LLM Training with JAX
- [GTC 2025] Scaling Transformers: Navigating Challenges and Innovations in Long-Context Modeling
Container images
Frameworks and supported models
We support and test the following JAX-based frameworks and model architectures. More details about each model and available containers can be found in their respective READMEs.
Nightly image builds use the latest JAX and XLA and are published to the ghcr.io/nvidia/jax image registry. The plain image tag ghcr.io/nvidia/jax:jax or ghcr.io/nvidia/jax:maxtext is simply the latest image. Append a date suffix to any image tag to specify a particular nightly build, e.g ghcr.io/nvidia/jax:jax-2026-06-01.
Stable builds are released at the end of each month to NVIDIA's NGC Catalog at the nvcr.io/nvidia/jax image registry. NGC images pin specific JAX releases. See NGC JAX release notes for more details.
| Framework | Models | Use cases | Container image (nightly build) | Container image (NGC release) |
|---|---|---|---|---|
| jax | pre-training | ghcr.io/nvidia/jax:jax | nvcr.io/nvidia/jax:YY.MM-py3 | |
| maxtext | DeepSeek V3.2, Gemma 4, Llama 4, Mixtral | pre-training | ghcr.io/nvidia/jax:maxtext | nvcr.io/nvidia/jax:YY.MM-maxtext-py3 |
| axlearn | Fuji | pre-training | ghcr.io/nvidia/jax:axlearn | |
| alphafold3 | Evoformer | inference | ghcr.io/nvidia/jax:alphafold |
Nightly 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 addition to the public CI, we also run internal CI nightlies on GB300, B300 SXM6, GB200, B200, DGX Spark, RTX PRO 6000 Blackwell, Jetson AGX Thor, H100 SXM 80GB, A100 SXM 80GB.
Experimental projects
- JAX-vLLM Rollout Offloading Bridge
- This project couples JAX training with vLLM inference to accelerate reinforcement-learning (RL) post-training