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:6fe367b373829722e362d466c965b86003bb624a148fa05a76c147b6628b49a3	RUN	`URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs sed -i "s@ f'jaxlib >={_minimum_jaxlib_version}, <={_jax_version}',@ f'jaxlib>=0.5.0',@" /opt/jax/setup.py`	08/08/2025 6:53 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_JAX_VERSION=25.08`	08/08/2025 6:53 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`NVIDIA_PRODUCT_NAME=JAX`	08/08/2025 6:53 PM UTC
sha256:65cea4695e603b1214fb3c751f5c20f6618b0b93131b8439d336edc8834b745f	RUN	`URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh`	08/08/2025 6:53 PM UTC
sha256:dabba5924a34d7a5c8a58fde5c19ba4bfb490212e9f155237c3eeb512ca1ab02	RUN	`RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 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`	08/08/2025 6:48 PM UTC
sha256:23b4c23e2d3296c2e9403dbf4a7f1f7539fea3c3370142ae0ef169ab5879ab8d	COPY	`/opt/transformer-engine /opt/transformer-engine`	08/08/2025 6:48 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4	ENV	`SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine`	08/08/2025 6:48 PM UTC
sha256:bdadb1f360866cb5c02557c76bec90d9ac4247761368352f7575b508f06b4199	RUN	`RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} sed -i 's/orbax-checkpoint/orbax-checkpoint==0.11.19/' ${SRC_PATH_FLAX}/pyproject.toml echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit`	08/08/2025 6:48 PM UTC
sha256:0c6758fbecc06b58963fbe3cc672029dc8cf9310982045ed605563225a4fec8b	RUN	`RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex for component in $(ls ${BUILD_PATH_JAXLIB}); do echo "-e file://${BUILD_PATH_JAXLIB}/${component}" >> /opt/pip-tools.d/requirements-jax.in; done echo "-e file://${SRC_PATH_JAX}[k8s]" >> /opt/pip-tools.d/requirements-jax.in EOF # buildkit`	08/08/2025 6:48 PM UTC
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1	RUN	`URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d`	08/08/2025 6:48 PM UTC