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`	04/16/2025 9:50 PM UTC
sha256:59f1ffe006963baf1c24e22d65c3900c7911cd5c7306fab7d7579dcdea80a822	RUN	`NVIDIA_BUILD_ID=159526777 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb SRC_PATH_MAXTEXT=/opt/maxtext pip install "werkzeug>=3.0.3"`	04/16/2025 9:50 PM UTC
sha256:b1cb2010f80af6261709be63b779653d241967278150012a2e7e1833a5467db2	RUN	`NVIDIA_BUILD_ID=159526777 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh`	04/16/2025 9:50 PM UTC
sha256:9c30d93a32ca63f1474f040a04701a6aca2d5b4432d168ed9e56ae1979e66409	ADD	`test-maxtext.sh /usr/local/bin`	04/16/2025 9:47 PM UTC
sha256:d144878496e4e06314f5677c054e3a96956fce96f783d5c413bf5a1d9e9886c8	RUN	RUN \|3 NVIDIA_BUILD_ID=159526777 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} echo "-r ${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in # specify some restrictions to speed up the build and # avoid pip to download and check all available versions of packages for pattern in \ "s\|absl-py\|absl-py>=2.1.0\|g" \ "s\|protobuf==3.20.3\|protobuf>=3.19.0\|g" \ "s\|tensorflow-datasets\|tensorflow-datasets>=4.8.0\|g" \ "s\|sentencepiece==0.1.97\|sentencepiece>=0.2\|g" \ "s\|tensorflow-text>=2.13.0\|tensorflow-text==2.18.1\|g" \ "s\|tensorflow>=2.13.0\|tensorflow==2.18.1\|g" \ ; do sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt; done # add new line in case requirements.txt does not end with a new line echo >> ${SRC_PATH_MAXTEXT}/requirements.txt for requirement in \ "tensorflow-metadata>=1.15.0" \ "seqio@git+https://github.com/google/seqio.git" \ ; do echo "${requirement}" >> ${SRC_PATH_MAXTEXT}/requirements.txt done EOF # buildkit	04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_MAXTEXT=/opt/maxtext`	04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb`	04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_BUILD_ID=159526777`	04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`NVIDIA_BUILD_ID=159526777`	04/16/2025 9:47 PM UTC
sha256:b0c7621c7577f60feee0d061232b6c8cd6d382db48480c38449313ab0deda8e7	RUN	`RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-04-16 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/.whl echo "transformer-engine @ file://$(ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/.whl)" > /opt/pip-tools.d/requirements-te.in EOF # buildkit`	04/16/2025 9:40 PM UTC