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:3e96afa6de3fb2cf4054284954823c5a3ebc04ab048631545f04e627e5d68047	COPY	`/opt/xla-source/.git /opt/xla-source/.git`	08/25/2023 10:01 PM UTC
sha256:de2eda6316c7f3f77eb12313edafd6671c628d854123bdd0cd7f052929bd22f5	COPY	`/opt/jax-source/.git /opt/jax-source/.git`	08/25/2023 10:01 PM UTC
sha256:9eb64320937486c5d8b6dd568feebb55f4852082207268a9e4c0b0032121540c	ADD	`build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/`	08/25/2023 10:01 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_XLA`	08/25/2023 10:01 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ARG	`SRC_PATH_JAX`	08/25/2023 10:01 PM UTC
sha256:8cdf9d7f23c2c7973646d7a17e00bd43568290a46b1632a15d606312fd8779f4	RUN	`SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-08-25 { pip install flax \|\| true; } && rm -rf ~/.cache/pip`	08/25/2023 10:01 PM UTC
sha256:30bcf085181c30c91037920e58a1d3f472ad4b9cdc7d0d8e27a9eeed1f428312	RUN	`SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-08-25 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/`	08/25/2023 10:01 PM UTC
sha256:ed5b2911ec6ae2999dd20a36b17420293256fde286d02c2cc499a437d2baaab1	COPY	`/opt/xla-source-no-git /opt/xla-source`	08/25/2023 10:01 PM UTC
sha256:3517e6a9f121ffc4e05e5d2bca6c0d953716b79e6f16c4352dd824c9690641f9	COPY	`/opt/jax-source-no-git /opt/jax-source`	08/25/2023 10:01 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NCCL_IB_SL=1`	08/25/2023 10:01 PM UTC