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.07`	07/25/2026 3:18 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_PRODUCT_NAME=JAX`	07/25/2026 3:18 AM UTC
sha256:cdeca2f85d84920407f27d6f9832300cf696a4b391febf8723bfc945056c4bfa	RUN	`URLREF_FLAX=https://github.com/google/flax.git#5768bb5902ffc01e578ba67a21af48d8457e6905 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-07-24 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c pip-finalize.sh`	07/25/2026 3:18 AM UTC
sha256:1e473df300bd2023303d928e6076dd6ca0f17bc58b98ba8d2163d43c43cf9e8c	RUN	`URLREF_FLAX=https://github.com/google/flax.git#5768bb5902ffc01e578ba67a21af48d8457e6905 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-07-24 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in`	07/25/2026 3:16 AM UTC
sha256:f130034fdf6e78064fed1f689fdeb0d27d9254620a5fee57fda24ad01b568c50	RUN	`URLREF_FLAX=https://github.com/google/flax.git#5768bb5902ffc01e578ba67a21af48d8457e6905 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-07-24 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}`	07/25/2026 3:16 AM UTC
sha256:9803051bf6bca464de8625ad37b49ef365658f1d4558234867e1d449b997ce19	ADD	`install-nsys-jax.sh /usr/local/bin`	07/25/2026 3:16 AM UTC
sha256:3e131e89cc4d9c30978580e7edf7edd58b70b79f84e96b5a5b16d5b923bda0a7	COPY	`nsys_jax/ /opt/nsys_jax`	07/25/2026 3:16 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`SRC_PATH_NSYS_JAX=/opt/nsys_jax`	07/25/2026 3:16 AM UTC
sha256:d7a7f7d5584a8766bbeacfc798c6a8c254613a049403e58c34bdd3d7334fc24f	RUN	`URLREF_FLAX=https://github.com/google/flax.git#5768bb5902ffc01e578ba67a21af48d8457e6905 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-07-24 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c echo "nvidia-cutlass-dsl[cu13]" >> /opt/pip-tools.d/requirements-cutlass-dsl.in`	07/25/2026 3:16 AM UTC
sha256:974b1b061e465d4050ec1998aa9e27a9cfdc206785e4ac329698137f3d7034b9	RUN	RUN \|8 URLREF_FLAX=https://github.com/google/flax.git#5768bb5902ffc01e578ba67a21af48d8457e6905 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-07-24 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -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	07/25/2026 3:16 AM UTC