JAX | NVIDIA NGC

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.

Layer	Label		Created
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1	WORKDIR	`/opt/maxtext`	10/17/2025 5:38 AM UTC
sha256:e15213cd0e1e19ef34cc335c1cca0c28eca368d4a2cd27285a5d8563f058fae3	RUN	`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:1fb119e6c84edca9ccf13ec9a0f9527ccb5f9aaeb299494860d5a36aa2717c0f	COPY	`pinned/25.10-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt`	10/17/2025 5:36 AM UTC
sha256:dc1502f5d478abf24dfbd4c7148f18287fb66699a4b0b3a976e1ef68e5558465	ADD	`test-maxtext.sh /usr/local/bin`	10/17/2025 5:36 AM UTC
sha256:556dbdc9c3323ea9ad71bb3c9a967b3306901e4b21c4fa6fcb5f4ea62069cb4f	RUN	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:8008f44ff29f4e2e813ae5509ae1bfaf498ff718b39c4cf00c546ac32a663ff2	RUN	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:b1e826a886a9a6938bab91a61f92461c2b527c46636be2e39ec1b769268ec508	RUN	`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:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_JETSTREAM=/opt/jetstream`	10/17/2025 5:36 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_MAXTEXT=/opt/maxtext`	10/17/2025 5:36 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8`	10/17/2025 5:36 AM UTC