JAX Release 26.10

The NVIDIA JAX Release 26.10 is made up of two container images available on NGC: JAX and MaxText.

Contents of the JAX container

This container image contains the complete source for the following software:

• JAX: /opt/jax
• XLA: /opt/xla
• Flax:/opt/flax
• TransformerEngine:/opt/transformer-engine

The MaxText container image is based on the JAX container. Additionally, it includes:

• MaxText: /opt/maxtext

The JAX runtime package jaxlib is prebuilt and installed in the default Python environment (/usr/local/lib/python3.10/dist-packages/jaxlib) in the container image.

Versions of packages included in both of these containers:

• CUDA 13.4.1
  • Please refer to the CUDA DL 26.09 release notes section for the list of libraries inherited from the CUDA container.

• JAX jax-v0.11.2
• MaxText f178b07
• Flax 4f42ed0
• XLA 91888df
• Transformer Engine v2.20.2 (6ea2a74a)

Driver Requirements

Release 26.10 is based on CUDA 13.4.3. For comprehensive and up-to-date driver compatibility information, please refer to the following documentation:

• NVIDIA CUDA Compatibility Guide - Compatibility information between CUDA versions and driver releases.
• CUDA Toolkit Release Notes - Driver version requirements and compatibility matrices.
• NVIDIA Drivers Download - Latest NVIDIA drivers.

Key Features and Enhancements

• Updated cuDNN FrontEnd library to version 1.26.
• Shipped JAX FFI extension headers (collectives, record) in the JAX wheel.

MoE

• Added sqrtsoftplus scoring to fused MoE routing in JAX. (TransformerEngine).
• Enabled Expert Parallelism to borrow XLA's NCCL communicator when supported (TransformerEngine).
• Supported single-process multi-device execution in MoE block (TransformerEngine).

Attention

• Fixed JAX fused-attention backward output-gradient sharding, rotated THD metadata handling, and compound DP/FSDP test-data placement (TransformerEngine).

Precision

• Added Hopper BF16 grouped GEMM support in JAX (TransformerEngine).
• Fixed FP8 custom-gradient state overwrite in MaxText NNX training.

Collective and Communication

• Added a new CollectiveReduce HLO operation for GPU collectives.

Compiler and Runtime

• JAX upgraded to v0.11.1 (from v0.10.2) with CUDA 13.4 support.
• Fixed quadratic tracing behaviour in JAX for improved compilation performance

JAX Toolbox

The JAX Toolbox projects focus on achieving the best performance and convergence on NVIDIA Ampere, Hopper, and Blackwell architecture families and provide the latest deep learning models and scripts for training and fine-tuning. These examples are tested against a nightly CI as well as each NGC container release to ensure consistent accuracy and performance over time.

Nightly Containers

In addition to projects, JAX Toolbox includes nightly containers for libraries across the JAX ecosystem.

Known Issues

• None.