Bump to 26.06 base container.

Try to fix TE's build.

Author: olupton

.github/container/Dockerfile.base

@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1-labs
-ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04
+ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.06-cuda13.3-devel-ubuntu24.04
 ARG GIT_USER_NAME="JAX Toolbox"
 ARG GIT_USER_EMAIL=jax@nvidia.com
 ARG JAX_TOOLBOX_REF=main
@@ -133,3 +133,7 @@ ENV MANIFEST_FILE="/opt/manifest.d/manifest.yaml"
 ADD manifest.yaml create-distribution.sh bump.sh /opt/manifest.d/
 
 COPY patches/ /opt/manifest.d/patches/
+
+# Temporary workaround for TransformerEngine, which added a >=1.25.0 requirement
+# when no public cuda-dl-base container had that CUDNN_FRONTEND_VERSION.
+ENV CUDNN_FRONTEND_VERSION=1.25.0

.github/container/Dockerfile.jax

@@ -132,6 +132,7 @@ COPY --from=builder ${SRC_PATH_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE
 RUN <<"EOF" bash -ex
 ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl
 echo "transformer-engine @ file://$(ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl)" > /opt/pip-tools.d/requirements-te.in
+echo "nvidia-cudnn-frontend==${CUDNN_FRONTEND_VERSION}" >> /opt/pip-tools.d/requirements-te.in
 EOF
 
 ## nvidia-cutlass-dsl[cu13]

.github/container/build-te.sh

@@ -138,13 +138,15 @@ pushd ${SRC_PATH_TE}
 # JAX, or the wheel-based installation of CUDA. Note that when we build TE as
 # part of building the JAX containers, JAX and XLA are not yet installed.
 python - << EOF
-import subprocess, sys, tomllib
+import os, subprocess, sys, tomllib
 with open("pyproject.toml", "rb") as ifile:
     data = tomllib.load(ifile)
 subprocess.run(
     [sys.executable, "-m", "pip", "install"]
     + [r for r in data["build-system"]["requires"]
-       if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")])
+       if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")]
+    + [f"nvidia-cudnn-frontend=={os.environ['CUDNN_FRONTEND_VERSION']}"]
+)
 EOF
 if [[ "${CCACHE}" == "1" ]]; then
     # Install ccache if not present (needs >= 4.1 for Redis remote storage support)

README.md

@@ -222,6 +222,7 @@ There are various other XLA flags users can set to improve performance. XLA flag
 
 | First nightly with new base container | Base container |
 | ------------------------------------- | -------------- |
+| 2026-07-08 | nvcr.io/nvidia/cuda-dl-base:26.06-cuda13.3-devel-ubuntu24.04 |
 | 2026-06-26 | nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 |
 | 2026-05-01 | nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 |
 | 2026-03-12 | nvcr.io/nvidia/cuda-dl-base:26.02-cuda13.1-devel-ubuntu24.04 |