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
01/21/2026 9:53 PM UTC
sha256:b54607517544aa2bbbdf2617efe01df7b7172d513116666ba4268636fb66b3cfRUN
NVIDIA_BUILD_ID=42221313 URLREF_MAXTEXT=https://github.com/google/maxtext.git#aca5b24931ebcbadb55a82e56ebffe8024874028 SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
01/21/2026 9:53 PM UTC
sha256:3a4f210a131863fc1d81da72f2d07668a9257a96fa06eaa829671333e58d9cc3COPY
pinned/26.01-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
01/21/2026 9:51 PM UTC
sha256:03e1922d03d1b8152316c1d9a2eaf2b6d325826947feba60fab663d8ba4a7475ADD
test-maxtext.sh /usr/local/bin
01/21/2026 9:51 PM UTC
sha256:88179cbead9055aa30efbd0b884ed44ee0e18c5788092e5390a0c51071e89b75RUN
RUN |3 NVIDIA_BUILD_ID=42221313 URLREF_MAXTEXT=https://github.com/google/maxtext.git#aca5b24931ebcbadb55a82e56ebffe8024874028 SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail for pattern in \
"s|^tensorflow$|tensorflow==2.18.1|g" \
"s|^tensorflow-text$|tensorflow-text==2.18.1|g" \
; do # tensorflow-cpu,tensorboard,tensorflow-text>=2.19.0 is incompatible with tensorflow==2.18.1 sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/dependencies/requirements/base_requirements/requirements.txt done # 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
01/21/2026 9:51 PM UTC
sha256:ce0067a87b00332cb8c5cb52100ebe28da38df6cfddaa54c8783c58b9bca369bRUN
RUN |3 NVIDIA_BUILD_ID=42221313 URLREF_MAXTEXT=https://github.com/google/maxtext.git#aca5b24931ebcbadb55a82e56ebffe8024874028 SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail echo "-e file://${SRC_PATH_MAXTEXT}" >> /opt/pip-tools.d/requirements-maxtext.in echo "-r ${SRC_PATH_MAXTEXT}/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
01/21/2026 9:51 PM UTC
sha256:e8ca200e0e78c3c251ea0b0f2003931200c6493aab0330964214967310b4b47aRUN
RUN |3 NVIDIA_BUILD_ID=42221313 URLREF_MAXTEXT=https://github.com/google/maxtext.git#aca5b24931ebcbadb55a82e56ebffe8024874028 SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
01/21/2026 9:51 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
01/21/2026 9:51 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#aca5b24931ebcbadb55a82e56ebffe8024874028
01/21/2026 9:51 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=42221313
01/21/2026 9:51 PM 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.