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

• Tutorials: Overview of JAX on GPUs, customizing JAX with NVIDIA CUDA Toolkit and going beyond what’s possible in native JAX.

• Guides: 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 JAX 26.06 container image released
  • Added support for Kimi K2-Thinking, K2.5, and K2.6 (text) models, including checkpoint conversion scripts.
  • Added Qwen3-30B-A3B-Base tokenizer and Qwen3.5 text-only decoder layer.
  • Added OLMo 3 7B/32B HuggingFace configs, stage-1 pretraining scripts, and numpy pretrain data pipeline.
  • Extended Gemma4 support: HuggingFace checkpoint conversion, vLLM adapter, layer-wise unit tests, MoE inference performance improvements, and multimodal evaluation (ChartQA).
  • Implemented DeepSeek, Gemma3, and Llama4 decoder layers in NNX.
• NGC JAX 26.05 container image released
  • performant jax.ragged_dot; XLA’s nvfp4 kernel codegen improvements; multi-stream collective support in XLA; PDL (programmatic dependent launch) enablement in XLA and improved D2H & H2D copy overlap with compute
• NGC JAX 26.04 container image released
  • 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

• [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

Tutorials

• JAX on GPU stack
• Writing High-Performance CuTe DSL kernels in JAX

Guides

• Tips for High-Performance LLMs with JAX on GPUs
• AutoPGLE on GPU workflows
• Profiling JAX programs on GPU
  • nsys-jax
• Resilient Training with Ray and JAX
• User Guide on Native XLA-FP8

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, Qwen3, Qwen 3.5, Mixtral, Kimi K2.6	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