Skip to main content
NVIDIA
JAX
Container
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
1-10 of 154 items

Modal Content

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.