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
7760b2131bf449b63e6ea75605894c0dd5a1d33f4f96354a6fa25369dfe28fd3CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
02/26/2026 1:21 AM UTC
215427b32f92408314e5d3f2489fe0b7c7054cab628820044f2d7641d6a94a7cRUN
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
ea9e7cd81f0474c22b2539337a75dc31a20b49a821380c504677431e71915048COPY
pinned/26.02-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
02/26/2026 1:16 AM UTC
7177e910d15ff964e8417f5b0f144665d0a96879bf3489fcfd71b51f8db843fcADD
test-maxtext.sh /usr/local/bin
02/26/2026 1:16 AM UTC
afca14f40a6a40326ef56db4926783df952b000cdb6677b009fcd28c9a1f6ab4RUN
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
61868d6e7ab67f4e36f95d2e27989053e3885e0e4ffef4a7781a8bea725c8b18RUN
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
25b0d47737d5177728584f2dc6ef7c3afe5f2eedda78e8970a0f9733d07ce55bRUN
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
6df67bd95996606ae1f146ae6c521f1b025bace1a7098388309e9a704ce679e5ARG
SRC_PATH_MAXTEXT=/opt/maxtext
02/26/2026 1:16 AM UTC
1ecf29da71843bbe39835afe63aa200f063df17e155cbb0a9c2debc75f9505c0ARG
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
85c8d6e07d431076196afbc94577331bb0b1b9661eca82f077ff434764c4e090ENV
NVIDIA_BUILD_ID=44841511
02/26/2026 1:16 AM UTC
...