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
b76f84525c9b6053f1df5c45c558a77e04119cbdb1ff2db5093cf8e113b2d396CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
03/24/2026 8:40 AM UTC
53b8a6094bf98ee3b4c09328d5d2c33020ba5508b8fd820747792154c3378cc4RUN
NVIDIA_BUILD_ID=46857070 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_oFKudKseHFbWh7aep1MY@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
03/24/2026 8:40 AM UTC
a51d8cf6ec1b91f79dcec7c2904591e2b4cbb15c492c34c83107f39a2fd3543bCOPY
pinned/26.03-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
03/24/2026 8:39 AM UTC
4000b2091ab1d17a46448219e03688125ff8de707b543aa3029eccceb7905261ADD
test-maxtext.sh /usr/local/bin
03/24/2026 8:39 AM UTC
16b48224cb18781365689d5faed3fa0112f7367d8492b719a981f4ec52130085RUN
RUN |3 NVIDIA_BUILD_ID=46857070 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_oFKudKseHFbWh7aep1MY@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel 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
03/24/2026 8:39 AM UTC
722968da24c2b185aa6944d877e3297621206232b3a753152e0dc323aac15d85RUN
RUN |3 NVIDIA_BUILD_ID=46857070 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_oFKudKseHFbWh7aep1MY@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel 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
03/24/2026 8:39 AM UTC
58bfb2755e099f954b66e718b9a4bf6bfe6cf2c74e59c17567b22fe0d1f3da1fRUN
RUN |3 NVIDIA_BUILD_ID=46857070 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_oFKudKseHFbWh7aep1MY@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
03/24/2026 8:39 AM UTC
92528fb9ba75cb151b671e4deb1831938ef8d31332f74e0e04fd3811670add69ARG
SRC_PATH_MAXTEXT=/opt/maxtext
03/24/2026 8:39 AM UTC
450d1c9986c697c64691d51b8c4e37d86af37a48cfe790e479cab2473da23700ARG
URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_oFKudKseHFbWh7aep1MY@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel
03/24/2026 8:39 AM UTC
a9d7031c99cc03f43037f45d2b2fdbc5f9d06641de4f7e2968e86f4057a67c45ENV
NVIDIA_BUILD_ID=46857070
03/24/2026 8:39 AM UTC
...