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
1230a1d72d1ed3121a84e5fd80d9fed37c85effedbaa524db5845bc227698858CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
04/23/2026 5:15 PM UTC
7aa845b698ba5d56232789b9d770856e288c4a16f2b252fc0606ee67f1235189RUN
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
a7061c2d34b174e1ce0a024a4ea84ee8ba6d511ed535d38d9f26af11d9360e6dCOPY
pinned/26.04-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
04/23/2026 5:14 PM UTC
5d41873a054e7217079d69751b2373c7a22c44b48566973aa9df694c7c81b575ADD
test-maxtext.sh /usr/local/bin
04/23/2026 5:14 PM UTC
878deff06baec805fdd6281069722c93efb434ca356bc950a5a92b303da414bdRUN
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
5aa3d1cc257615a114fd7106edfb8c4ec82d3848b975bc943b3815fd4fd2ea4dRUN
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
43d1b73647b687e859699d3b06a79329965d77ee7ebb60a20f83be429e400b6cRUN
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
5f3f71dc4d511720c078fdbd7ef954279e3046ed3be46a4ad6003034b6050e53ARG
SRC_PATH_MAXTEXT=/opt/maxtext
04/23/2026 5:14 PM UTC
ea208966a176b2da87f67e0972a54987e93637ac797ce266ee6ffe855b798895ARG
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
d82c0c36a8659ac27b60e756fb2c0b4abf78fa2f31256c43cd6055c45cefdebbENV
NVIDIA_BUILD_ID=49311194
04/23/2026 5:14 PM 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.