ContainersNVIDIADeepSeek-R1NVIDIALlama-3.1-Nemotron-70B-InstructNVIDIAPyTorch
CollectionsNVIDIAOmniverse Kit (FB)NVIDIADeepStream SDKNVIDIAOmniverse Kit App Streaming
ModelsNVIDIAStyleGAN3 pretrained modelsNVIDIAPeopleNetNVIDIATrafficCamNet
ResourcesNVIDIARiva Skills Quick Start
Helm ChartsNVIDIAGPU OperatorNVIDIANVIDIA NIM Operator
NVIDIA
NVIDIA
JAX
Container
NVIDIA
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.

LayerLabelCreated
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_JAX_VERSION=26.02
02/26/2026 12:30 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
02/26/2026 12:30 AM UTC
sha256:d4446458e23260233f740b2240456b015103820c1729c7e96a607dc25820e70dRUN
URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-02-25 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
02/26/2026 12:30 AM UTC
sha256:1ae5feaaf155ef8a3743b063c6fa24be1f3e4140144c9ceb52981b9b0fc2bc67COPY
pinned/26.02-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
02/26/2026 12:30 AM UTC
sha256:baa96ad34f98382edca7d926217620a24b3ff6b526655f1ef458a7f2d485e472RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-02-25 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
02/26/2026 12:28 AM UTC
sha256:003f7e82157ab63dde3e4076ed3d066821489e2bea7308658895c97b3e9ba8f0COPY
/opt/transformer-engine /opt/transformer-engine
02/26/2026 12:28 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
02/26/2026 12:28 AM UTC
sha256:24d516e3351bf330705c894577c0ad8adbb47f9ece637120e934b186b08e3562RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-02-25 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit
02/26/2026 12:28 AM UTC
sha256:051aa59b5b9f9b44048f9f440d9d6368ca636823b8ec642984c2e23d085cf3d1RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-02-25 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
02/26/2026 12:28 AM UTC
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1RUN
URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-02-25 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d
02/26/2026 12:28 AM UTC
...