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:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_JAX_VERSION=26.04`	04/23/2026 5:10 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_PRODUCT_NAME=JAX`	04/23/2026 5:10 PM UTC
sha256:c6e4b36e806d1847cff318015c1e1377016b382aed4b907d569bd333858c78fe	RUN	`URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh`	04/23/2026 5:10 PM UTC
sha256:d10f98e3ef3838fba106ff77cd5977c0201cf45d84c8820fe4929e79fdc057ee	COPY	`pinned/26.04-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt`	04/23/2026 5:09 PM UTC
sha256:2496fcc20b6e04e19febe83d5cda0af20bdf9da25e6092c14eceaeac3243bdec	RUN	`URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in`	04/23/2026 5:08 PM UTC
sha256:1061fc49a5178e4329bc472cc5b2a76af94685dde77563fd25255c5cac59410f	RUN	`URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}`	04/23/2026 5:08 PM UTC
sha256:3514e9c393440bd66082303e321cf586beb9aae5b358aa1442d614dc1d2fda09	ADD	`install-nsys-jax.sh /usr/local/bin`	04/23/2026 5:08 PM UTC
sha256:13336517bd202f06e7fe50a42f9407206d5897ddf65e9ff26b000faae8bb4f74	COPY	`nsys_jax/ /opt/nsys_jax`	04/23/2026 5:08 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`SRC_PATH_NSYS_JAX=/opt/nsys_jax`	04/23/2026 5:08 PM UTC
sha256:b1efd299f774981adb32a6f5cf12402c194bcc151c33471838f6f53ea0b6708e	RUN	RUN \|8 URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 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/23/2026 5:08 PM UTC