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
02/26/2026 1:21 AM UTC
sha256:6713f648e8097dfc22da1f654c533c50427a36e5a54b2f9ee11e6863ae4fb9bbRUN
NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
02/26/2026 1:21 AM UTC
sha256:54bd4758e2aea5512b7d6192376dacfc0fd1fcf64bc4e3987492e81754e04d07COPY
pinned/26.02-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
02/26/2026 1:16 AM UTC
sha256:526474a8ff5a5892f3b454ec51a0d87ed1b86bab1328a745349580a094f545cfADD
test-maxtext.sh /usr/local/bin
02/26/2026 1:16 AM UTC
sha256:a6b328f241ffa022e32187136a56aaecbe24ee17bd63f9428042e0656600b016RUN
RUN |3 NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@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
02/26/2026 1:16 AM UTC
sha256:d241a2f602c6b04132e8d6b59aaa7369bd38c2f9f510ba231cd2baf181e96062RUN
RUN |3 NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@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
02/26/2026 1:16 AM UTC
sha256:47966881c48ee8335283c790dfa6b6ae49ae3df297908a0fce3643be9a7764e3RUN
RUN |3 NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@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
02/26/2026 1:16 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
02/26/2026 1:16 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel
02/26/2026 1:16 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=44841511
02/26/2026 1:16 AM 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.