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
3294dc5f5d9167a03e737e060152e42402af20669a5f9526a3e5220f195c6a7b	CONFIG	`Entrypoint /opt/nvidia/nvidia_entrypoint.sh`	01/21/2026 9:39 PM UTC
c6aa6ab5247aa1149ea3fe826b440ecb067e4fdbd7b978eaa49a9c1341f2c0f3	ENV	`NVIDIA_PRODUCT_NAME=JAX`	01/21/2026 9:39 PM UTC
0574b0741f975705cb2cdb981a3f3448b429e0a1855af895648026df98fde426	RUN	`URLREF_FLAX=https://github.com/google/flax.git#v0.12.1 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-01-21 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh`	01/21/2026 9:39 PM UTC
536cca21db3d453a033c034ae3be6972fec1ddcdfe262d551e62cf314449af0e	COPY	`pinned/26.01-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt`	01/21/2026 9:38 PM UTC
48b49044afc5a310b14891f2644fec6e7c5b62b8dcb7b3a9caf173fa9157aee6	RUN	`RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.12.1 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-01-21 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`	01/21/2026 9:35 PM UTC
7f9139a4de08b5f5022d0c8e81c9b9a506f7bfc701a8384e8ed1b03ae01a2921	COPY	`/opt/transformer-engine /opt/transformer-engine`	01/21/2026 9:35 PM UTC
7ff33d446e371d63a4366c9a0c3b6270f71aa5fb8169ad5857defa346877ef5c	ENV	`SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine`	01/21/2026 9:35 PM UTC
038774c553924774f50ef36468ab1097b24485b6e8f370a52c2b196c7c5541b0	RUN	`RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.12.1 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-01-21 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} sed -i 's/jax>=0\.8\.1/jax/' ${SRC_PATH_FLAX}/pyproject.toml echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit`	01/21/2026 9:35 PM UTC
07f362b022ec8367abd884f3f73a24fb74d5ec46687b9f91bc9ae6be5e56a72a	RUN	RUN \|7 URLREF_FLAX=https://github.com/google/flax.git#v0.12.1 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-01-21 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 echo "kubernetes @ git+https://github.com/kubernetes-client/python.git#release-34.0" >> /opt/pip-tools.d/requirements-jax.in echo "urllib3>2.6.0" >> /opt/pip-tools.d/requirements-jax.in EOF # buildkit	01/21/2026 9:35 PM UTC
574975175b6a39b5311d93a9ba94037cbecf54ef451e4a9a245775b0484c0b13	RUN	`URLREF_FLAX=https://github.com/google/flax.git#v0.12.1 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-01-21 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d`	01/21/2026 9:35 PM UTC