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:d5b0f35822dd3a6d7dd6ea528a7800b0448aa757498e3801bcb8c58f7484f07b	COPY	`/opt/xla-source/.git /opt/xla-source/.git`	10/20/2023 8:42 PM UTC
sha256:d4fe8e30eed895dde8c7edec8842a053328e3fc7a4b6a2947363ec03d2604cf4	COPY	`/opt/jax-source/.git /opt/jax-source/.git`	10/20/2023 8:42 PM UTC
sha256:b8cf68eafb357b0345781e9467d357c4161354ba53c593034201d94e62b585b7	ADD	`build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/`	10/20/2023 8:42 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_XLA`	10/20/2023 8:42 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_JAX`	10/20/2023 8:42 PM UTC
sha256:4817dfb84baf486370d860e8d30d92e6cf56f325d3a6599cfef74c2f2d21e682	RUN	`SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-10-20 { pip install flax \|\| true; } && rm -rf ~/.cache/*`	10/20/2023 8:42 PM UTC
sha256:66290980827ac9b035755dcccd488cb2ab208ad4c45efd7ef18b3fff5d7cf4d9	RUN	`SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-10-20 pip --disable-pip-version-check install ${SRC_PATH_JAX}/dist/*.whl && pip --disable-pip-version-check install -e ${SRC_PATH_JAX} && rm -rf ~/.cache/pip/`	10/20/2023 8:42 PM UTC
sha256:92dc1befe7e890e3f0256bf9831f29bfd49039396c51867e269cac5bb97fe397	COPY	`/opt/xla-source-no-git /opt/xla-source`	10/20/2023 8:42 PM UTC
sha256:25d689a19b52e6518ee7ccc24090ead16c3940021f68401e5b6186103aeb3717	COPY	`/opt/jax-source-no-git /opt/jax-source`	10/20/2023 8:42 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NCCL_NVLS_ENABLE=0`	10/20/2023 8:42 PM UTC