Introduction
A public CI, Docker images, and optimized examples for JAX on NVIDIA GPUs
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.
Where to go next
- Frameworks & Supported Models — the frameworks and model architectures we support and test, and their containers.
- Build Pipeline Status — live build and test status for every published container.
- Environment Variables — the XLA and NCCL flags baked into the images for performance tuning.
- Profiling — how to profile JAX programs on GPU.
- Container Versions — base-container history and staging releases.
Containers
In all cases, ghcr.io/nvidia/jax:XXX points to the 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, and A100 SXM 80GB.
License
JAX Toolbox is released under the Apache 2.0 License.