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
10/17/2025 5:38 AM UTC
sha256:e15213cd0e1e19ef34cc335c1cca0c28eca368d4a2cd27285a5d8563f058fae3RUN
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
sha256:1fb119e6c84edca9ccf13ec9a0f9527ccb5f9aaeb299494860d5a36aa2717c0fCOPY
pinned/25.10-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
10/17/2025 5:36 AM UTC
sha256:dc1502f5d478abf24dfbd4c7148f18287fb66699a4b0b3a976e1ef68e5558465ADD
test-maxtext.sh /usr/local/bin
10/17/2025 5:36 AM UTC
sha256:556dbdc9c3323ea9ad71bb3c9a967b3306901e4b21c4fa6fcb5f4ea62069cb4fRUN
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
sha256:8008f44ff29f4e2e813ae5509ae1bfaf498ff718b39c4cf00c546ac32a663ff2RUN
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
sha256:b1e826a886a9a6938bab91a61f92461c2b527c46636be2e39ec1b769268ec508RUN
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
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_JETSTREAM=/opt/jetstream
10/17/2025 5:36 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
10/17/2025 5:36 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8
10/17/2025 5:36 AM UTC
...