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
2b96a5606af2ba950622de28e9c3c8a8896ac2e02a2f7b3721a5182722e420efCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
01/21/2026 9:53 PM UTC
4b8d3a7f5734fa152821b92f2d2ecffb2478fea05268dc37c313387ce28e3985RUN
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
58b91d0d86ae87aeb2cd87fd687a9b9cb120d7a12e4eb1249d29297cacf0b77bCOPY
pinned/26.01-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
01/21/2026 9:51 PM UTC
16f450b9274b043070934b0baba55033b7d4700399dcecb722636c3d37e9c973ADD
test-maxtext.sh /usr/local/bin
01/21/2026 9:51 PM UTC
e0c736068dd5e814826a00fd10cda8fb7a0238f3cb4f7c36bfefd2ce2b5b4a70RUN
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
67e21b8c6ca50a832c77f6c52cedc88b6a3182ec9c769bb25525fa5399526472RUN
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
c68d0b6c62e23bb504aa15c4f14e373fc9db11e6dc1df47276321357afc6245cRUN
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
e3eb8dcf48d5f79c00f9cd8dd8a55b3ac52db7e8f07e58f8b0fb7955bbff3a52ARG
SRC_PATH_MAXTEXT=/opt/maxtext
01/21/2026 9:51 PM UTC
37ce53d4b62712bb2acf4263dea781bec99ef849a4093f96c075426280c2cfd0ARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#aca5b24931ebcbadb55a82e56ebffe8024874028
01/21/2026 9:51 PM UTC
26c3b9fa8d103b51a127a3e78b5849b55e1757c5187636bd16f971183c22f544ENV
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.