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.05`	05/21/2026 10:50 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_PRODUCT_NAME=JAX`	05/21/2026 10:50 AM UTC
sha256:932ab85622293720639c690caaa4a30460e6b200d3fa03260128f2d9a14550af	RUN	`URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh`	05/21/2026 10:50 AM UTC
sha256:e380d606fefb17da12f0fae4ae293687d6c49483b9fb84bed0c58930737f9472	COPY	`pinned/26.05-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt`	05/21/2026 10:49 AM UTC
sha256:0a88b5dbfda400b4505fc266d771ecfc0f71df1878426697ed5089b8a05291cd	RUN	`URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in`	05/21/2026 10:48 AM UTC
sha256:813facc106ea9c197ffa61ec424e8b8307547d3ac5912381d1bff07247b295fe	RUN	`URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}`	05/21/2026 10:48 AM UTC
sha256:827b1ac4dca9e6baf6b7fc7cc7b151cb9779a6f2eb3912ba91ec48847dfeed55	ADD	`install-nsys-jax.sh /usr/local/bin`	05/21/2026 10:48 AM UTC
sha256:341e33fd1ab0295c1f7e9550213203bda9118c70a3587358144cd993aafdf7a4	COPY	`nsys_jax/ /opt/nsys_jax`	05/21/2026 10:48 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`SRC_PATH_NSYS_JAX=/opt/nsys_jax`	05/21/2026 10:48 AM UTC
sha256:001bb1d3a63507d7e04a97a40bf52ed995b7795c3c42296a6ae992f27c99d5f1	RUN	`URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "nvidia-cutlass-dsl[cu13]" >> /opt/pip-tools.d/requirements-cutlass-dsl.in`	05/21/2026 10:48 AM UTC