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`	05/08/2024 5:00 AM UTC
sha256:a7ce708dd4fbbc27700ebdc3ae2e167593f3a5cb5004ff8e2a6e8469f9d9fd6e	RUN	`pip-finalize.sh`	05/08/2024 5:00 AM UTC
sha256:8e1cd6717a6ed3c5ca2d39b63d08753c4956961e9fcec6e3d96e2a738747d7e3	ADD	`test-maxtext.sh /usr/local/bin`	05/08/2024 4:58 AM UTC
sha256:65f7827fc9ff39161c0ba426e4c574408f38c73ed1af7800cae4a455fdf4392e	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
sha256:c632f6fcb58231e2f90838bdbea90204011029ec64c8685f26593365a42c8ee5	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
sha256:186ed4e9b4515aff729b0bebc12a8722381361b76c51e3ebbe89ac48c0e56e86	ADD	`maxtext-mha.patch /opt`	05/08/2024 4:58 AM UTC
sha256:cc5620f21fecb845dc928e81839fe0fa659b96faa9eb8194703da63f25ee01bf	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
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_MAXTEXT`	05/08/2024 4:58 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_BUILD_ID=92034665`	05/08/2024 4:58 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`NVIDIA_BUILD_ID`	05/08/2024 4:58 AM UTC