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
04/23/2026 5:15 PM UTC
sha256:085dd8d9afc01334f534f697f1e943f8a3301376d31eb39e3b18b0c1a4c87391RUN
NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
04/23/2026 5:15 PM UTC
sha256:41e5181ca78c46738360314df044da6c0da28479e9ea04f6990d7d7244ba3aa5COPY
pinned/26.04-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
04/23/2026 5:14 PM UTC
sha256:9a91ef40d740741a9597558061180e0177c3eda3b1ee1c59482685bfe15bc1f3ADD
test-maxtext.sh /usr/local/bin
04/23/2026 5:14 PM UTC
sha256:9138061b14b96f79b13027e4167bd811c321fe1f859f4c9e7d2b88f4a427f23dRUN
RUN |3 NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -exuo pipefail # Avoid old tensorboard constraint from google-cloud-aiplatform[tensorboard]. sed -i "/^cloud-accelerator-diagnostics/d" \
${SRC_PATH_MAXTEXT}/src/dependencies/requirements/base_requirements/requirements.txt # 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
04/23/2026 5:14 PM UTC
sha256:a6c1412f14ee3cc1d533f6b677c2b8bea5b0500c8037c01cdb0829bf1d8dc72fRUN
RUN |3 NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-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}/src/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
04/23/2026 5:14 PM UTC
sha256:f19299bc21503897b83052481e9f7205f0135f42cd0437cae2daf7d37cd8adadRUN
RUN |3 NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
04/23/2026 5:14 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
04/23/2026 5:14 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel
04/23/2026 5:14 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=49311194
04/23/2026 5:14 PM UTC
...