Staging Containers

JAX-Toolbox staging containers host pending NVIDIA-authored XLA and JAX enhancements for NVIDIA GPU. These are pending PRs that are awaiting upstream review and merge in OSS OpenXLA and JAX repositories. For this initiative, we publish scale-training tagged containers — jax-scale-training.

The scale-training container is generated and published every other Saturday at 00:00 UTC. Refer to STAGING.md for more information on the underlying XLA staging branch, the pending PRs included, and the corresponding JAX commit used.

Staging releases

The underlying CUDA container used for building -scale-training containers can be seen from the corresponding dated version in Container Versions.

To check the versions of libraries in the container, run:

$
docker run --rm --gpus all ghcr.io/nvidia/jax:jax-scale-training-2026-04-24 -c '
>
echo "=== CUDA Toolkit ===" && echo ${CUDA_VERSION}
>
echo "=== cuDNN ===" && echo ${CUDNN_VERSION}
>
echo "=== NCCL ===" &&  echo ${NCCL_VERSION}
>
echo "=== Python Packages ===" && pip list | grep -iE "jax|flax|equinox|optax|chex|orbax|numpy|scipy|nvidia-"
>
'