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
03/24/2026 8:40 AM UTC
sha256:5017ba0cae97aacf0cd28b83f2fe9f27ea1e0ff5b0e07efdf67d09dde29534e9RUN
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
sha256:87f55e8320c5cd097a5c485b447d65238e002560e63ddc61710c7294d2cedc8aCOPY
pinned/26.03-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
03/24/2026 8:39 AM UTC
sha256:5361c818b3f6ac0a571cae17f3fd9d5df668df7fcabef1d64713fc7ef7791e0bADD
test-maxtext.sh /usr/local/bin
03/24/2026 8:39 AM UTC
sha256:f8fa48ffe6e62865dae9bbadf805054462546d42c00e2f23cdd1d8691af517fbRUN
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
sha256:76fc42e77ade5d6fe7f2d124a6ea395d57027f64b1631ea88f78a714ee0592c1RUN
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
sha256:6f24de8231bc56f5894d573e4320145a34a2a881e82d10c79c87b3fa13224a62RUN
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
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
03/24/2026 8:39 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
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
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=46857070
03/24/2026 8:39 AM UTC
...