Skip to main content
NVIDIA
JAX
Container
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
08/21/2026 6:29 PM UTC
sha256:6ac7e202900ebb11eed6f84b255e93f953fea690add4ac38a15edc762676ca51RUN
NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c pip-finalize.sh
08/21/2026 6:29 PM UTC
sha256:dd0375255cdd0d4a5a56d2a8af971d4f2aaaf7e9e47ec5bfc277bebb45450ce4COPY
pinned/26.08-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
08/21/2026 6:26 PM UTC
sha256:ac0df7e7a23df864d2b4eb2100e2565cc3ad4c84646054e0365fedef5c5ec052ADD
test-maxtext.sh /usr/local/bin
08/21/2026 6:26 PM UTC
sha256:f65589167301410224d8b2e280b89595d7f57d821ab8961b6f75549c8a9d2bd7RUN
RUN |3 NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada 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
08/21/2026 6:26 PM UTC
sha256:085f48e1907060031ec226bff0fc2dfb96bf71c2965a01c3d6ff40a0479619c8RUN
RUN |3 NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada 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
08/21/2026 6:26 PM UTC
sha256:f2461c8c6c6a179fc49d84abf9d51e4e62c93cc964994561cdcf6eb42db3128dRUN
RUN |3 NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
08/21/2026 6:26 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
08/21/2026 6:26 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada
08/21/2026 6:26 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=63903073
08/21/2026 6:26 PM UTC
1-10 of 156 items

Modal Content

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.