JAX Release 26.02
The NVIDIA JAX Release 26.02 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.1.U1
- Please refer to the CUDA DL 26.02 release notes section for the list of libraries inherited from the CUDA container.
Driver Requirements
Release 26.02 is based on CUDA 13.1 U1 (Toolkit) which requires CUDA Driver release 590.48. Please refer to the latest Drivers and CTKsupport table for additional information.
For a complete list of supported drivers, see the CUDA Application Compatibility topic. For more information, see CUDA Compatibility and Upgrades.
Key Features and Enhancements
This JAX release includes the following key features and enhancements.
- The current release is based on JAX 0.9.0.1 and CUDA 13.1 U1
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.
| Container | Type | Image URI |
|---|---|---|
| jax | - | ghcr.io/nvidia/jax:jax-YYYY-MM-DD |
| maxtext | LLM framework | ghcr.io/nvidia/jax:maxtext-YYYY-MM-DD |
| equinox | layer library | ghcr.io/nvidia/jax:equinox-YYYY-MM-DD |
| axlearn | LLM framework | ghcr.io/nvidia/jax:axlearn-YYYY-MM-DD |
Known Issues
-
The version of cuBLAS included in the container has a known issue where some cublasLtMatmul kernels can lead to incorrect results when executed concurrently with another kernel that uses Tensor Memory on GPUs with Compute Capability 10.x and 11.x.