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
3c57e0561a39ba790eddcc31ee79ce585ab4b6ee2929a6107f026eaf786fdaadCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
10/17/2025 5:32 AM UTC
d142d3f6dd143dfcaf5e3556243b2a0554430082f83cef10b54ceed7cdfae786ENV
NVIDIA_PRODUCT_NAME=JAX
10/17/2025 5:32 AM UTC
9bb33e17c522307a25a795cf8484423f2ce9851c7931c842d42fe4e6f4067ce1RUN
URLREF_FLAX=https://github.com/google/flax.git#v0.11.2 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-10-16 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
10/17/2025 5:32 AM UTC
d4881fee811f6ed2e54f7c42752dc6b6e5b50ed488d3edbc27a04758dd1c7e59COPY
pinned/25.10-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
10/17/2025 5:31 AM UTC
d98fe25fca4bec77ec167f58a482d8239013af4bb4428bed221f844c0c77c761RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.11.2 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-10-16 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
10/17/2025 5:30 AM UTC
131dbd3d16b3ab22d0255543610e0730b44dbf55697e4d63504baff38adb319dCOPY
/opt/transformer-engine /opt/transformer-engine
10/17/2025 5:30 AM UTC
3397e399253e4e08827a1043898761d2cb075ede5888f101c04a814da04d6482ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
10/17/2025 5:30 AM UTC
9bc93e84495d818366f073e886396c0166b4b6a492004cef856307dde3218f85RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.11.2 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-10-16 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
10/17/2025 5:30 AM UTC
4e9d7f100fd4e138c83b37b84336f752142b835ea87c256fc313f81bce4fbe16RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.11.2 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-10-16 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
10/17/2025 5:30 AM UTC
74ae55d1d7f98e2ab7654c94ded62aaaff703eeacb7ea9d15550fbf371aa6f32RUN
URLREF_FLAX=https://github.com/google/flax.git#v0.11.2 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-10-16 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d
10/17/2025 5:30 AM UTC
...