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
a6c05f3892586f0771c5492be331b7c73df95c42d932bc5ab7242a87b197b330CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
10/17/2025 5:38 AM UTC
16cd6ef245fdd60eb616da64e25e7bf54b4823f40fa1a21cfc4bda15129ca96cRUN
NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream pip-finalize.sh
10/17/2025 5:38 AM UTC
7c43374a6983d65fbe577b0050997a40a0090ec1d45f2956119d5ebefd852e94COPY
pinned/25.10-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
10/17/2025 5:36 AM UTC
8b71e913f9ee04b2676513695183fde056d8322d9491158b2772524da4364489ADD
test-maxtext.sh /usr/local/bin
10/17/2025 5:36 AM UTC
9557ac1635a255d55674ecf9894c435f48e17764f4e67a71fb5c835a71b8fde1RUN
RUN |5 NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail for pattern in \
"s|tensorflow>=2.19.1|tensorflow==2.18.1|g" \
"s|tensorboard>=2.19.0|tensorboard>=2.18,<2.19|g" \
"s|tensorflow-text>=2.19.0|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}/pyproject.toml done EOF # buildkit
10/17/2025 5:36 AM UTC
6518d3e4c12364936a011a93d540a3a40d192e52db31832de27c45a768442e2eRUN
RUN |5 NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /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/install_maxtext_extra_deps/extra_deps_from_github.txt" >> /opt/pip-tools.d/requirements-maxtext.in #echo "-e file://${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
10/17/2025 5:36 AM UTC
14a5bfbc9a5f6e11f5cc089465003c0553c8a0f5839b17f6f1339cbb90cd3896RUN
RUN |5 NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} EOF # buildkit
10/17/2025 5:36 AM UTC
5df4c97a763fe829fdd67b53d27956725d6d610b201e1697547ded95a8521364ARG
SRC_PATH_JETSTREAM=/opt/jetstream
10/17/2025 5:36 AM UTC
c33150bf4f1ec31a7e7db2d13699bd3842ab769f32fd041585d90f4175a3daa7ARG
SRC_PATH_MAXTEXT=/opt/maxtext
10/17/2025 5:36 AM UTC
cd007352a1b3483e5c414be6d70b2c49db56ee6bec26be31bf4055e4c27cdb11ARG
URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8
10/17/2025 5:36 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.