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
5bdf2ee4d9b829a594a2bfcf72fcf0013fb0770d84dbfd08144e36e71b4eeffc	CONFIG	`Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext`	05/08/2024 5:00 AM UTC
a0580e068d6a694581ed8e2a098e9806e27349f1b5d3f11043cbc8e7c41dc4af	RUN	`pip-finalize.sh`	05/08/2024 5:00 AM UTC
d6b617bc364c13584f17531f5d1f23c30795cf867361f73f03295463fac82eba	ADD	`test-maxtext.sh /usr/local/bin`	05/08/2024 4:58 AM UTC
59a1fc99556217b36497d371eae93df624e26c7d034aceefa3b795d0d1500a69	RUN	`RUN \|2 NVIDIA_BUILD_ID=92034665 SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex cat "${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit`	05/08/2024 4:58 AM UTC
6bb29b88ec17298ec1bfcf82b5bfc5b3d94ee069786f3ab515dbd16e275d31f5	RUN	`NVIDIA_BUILD_ID=92034665 SRC_PATH_MAXTEXT=/opt/maxtext cd "${SRC_PATH_MAXTEXT}" && patch -p1 < /opt/maxtext-mha.patch && git diff`	05/08/2024 4:58 AM UTC
3ffd0250b6397c07d0c42f8bfa1b635e84ac26d62ddd72e077f97fa1d2e3e76c	ADD	`maxtext-mha.patch /opt`	05/08/2024 4:58 AM UTC
8d05a7b73636e75a45be7fb5c06f9d820bfeb3f64b8442b088d9db1e664226eb	RUN	`RUN \|2 NVIDIA_BUILD_ID=92034665 SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex cat ${MANIFEST_FILE} get-source.sh -l maxtext -m ${MANIFEST_FILE} EOF # buildkit`	05/08/2024 4:58 AM UTC
8b76cd683606376fcb36ea416f930ba0ae707615c2888b6543c1d7311c0fd723	ARG	`SRC_PATH_MAXTEXT`	05/08/2024 4:58 AM UTC
594a58d19c93103001357829b8fb157872c2b8c869541bca8f54df8765d64a1c	ENV	`NVIDIA_BUILD_ID=92034665`	05/08/2024 4:58 AM UTC
59ea001a11c3812e579044f8b3616f530bd440d4d5977ff5bbe3bd73f07d14a7	ARG	`NVIDIA_BUILD_ID`	05/08/2024 4:58 AM UTC