NVIDIA
NVIDIA
JAX
Container
NVIDIA
NVIDIA
JAX

JAX is a framework for high-performance numerical computing and machine learning research. It includes Numpy-like APIs, automatic differentiation, XLA acceleration and simple primitives for scaling across GPUs and supports an ecosystem of libraries.

LayerLabelCreated
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
/opt/maxtext
07/25/2026 4:25 AM UTC
sha256:c45ceb7a453b172684e8d5f4bc570aec6404dd37b52847c38066b341a3503d98RUN
NVIDIA_BUILD_ID=59513657 URLREF_MAXTEXT=https://github.com/google/maxtext.git#454e55f4fedc4b11605378fc88d3e9afea353f11 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c pip-finalize.sh
07/25/2026 4:25 AM UTC
sha256:edce7e946f078349740e6b452f2b8b83a1439c6886401dc9266730cf2ff27194ADD
test-maxtext.sh /usr/local/bin
07/25/2026 4:22 AM UTC
sha256:808967feb756dc3b8f93166f52d805d6a1e423cb295c5775cb48821c5b819c57RUN
RUN |3 NVIDIA_BUILD_ID=59513657 URLREF_MAXTEXT=https://github.com/google/maxtext.git#454e55f4fedc4b11605378fc88d3e9afea353f11 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -exuo pipefail # Avoid old tensorboard constraint from google-cloud-aiplatform[tensorboard]. sed -i "/^cloud-accelerator-diagnostics/d" \
${SRC_PATH_MAXTEXT}/src/dependencies/requirements/base_requirements/requirements.txt # remove custom build hook that conflicts with our dependencies introduced in # https://github.com/AI-Hypercomputer/maxtext/commit/de886f9a918d17cfb025a353e7da2bc4d64c26b5 sed -i -e '/^\[tool\.hatch\.build\.targets\.wheel\.hooks\.custom\]$/d' \
-e '/^path = "build_hooks\.py"$/d' \
${SRC_PATH_MAXTEXT}/pyproject.toml EOF # buildkit
07/25/2026 4:22 AM UTC
sha256:e82740392c12a245d6f671c80128088c8a57c7ea2b288a4310ea549be12170c7RUN
RUN |3 NVIDIA_BUILD_ID=59513657 URLREF_MAXTEXT=https://github.com/google/maxtext.git#454e55f4fedc4b11605378fc88d3e9afea353f11 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -ex -o pipefail echo "-e file://${SRC_PATH_MAXTEXT}" >> /opt/pip-tools.d/requirements-maxtext.in echo "-r ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in # MaxText's base_requirements/requirements.txt (above) does not include google-jetstream,
# but inference code paths (maxtext.inference.maxengine, maxtext.input_pipeline.packing.prefill_packing)
# import the `jetstream` package at module load time and DECOUPLE_GCLOUD=TRUE does not
# unblock prefill_packing. Now re-add this dependency
grep -E '^google-jetstream' ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/requirements.txt \
  >> /opt/pip-tools.d/requirements-maxtext.in
EOF # buildkit
07/25/2026 4:22 AM UTC
sha256:6fa4459892305589ae1af7454b7fe02e39130a270a456762f53403f73e140a0cRUN
RUN |3 NVIDIA_BUILD_ID=59513657 URLREF_MAXTEXT=https://github.com/google/maxtext.git#454e55f4fedc4b11605378fc88d3e9afea353f11 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
07/25/2026 4:22 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
07/25/2026 4:22 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#454e55f4fedc4b11605378fc88d3e9afea353f11
07/25/2026 4:22 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=59513657
07/25/2026 4:22 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
NVIDIA_BUILD_ID=59513657
07/25/2026 4:22 AM UTC
...

NVIDIA uses cookies to improve your experience on our web site. We and our third-party partners also use cookies and other tools to collect and record information you provide as well as information about your interactions with our websites for performance improvement, analytics, and to assist in marketing efforts. By clicking "Accept All", you consent to our use of cookies and other tools as described in our Cookie Policy. You can manage your cookie settings by clicking on "Manage Settings." By continuing to use this site or by clicking one of the buttons below, you agree to our Terms of Service (which contains important waivers). Please see our Privacy Policy for more information on our privacy practices.