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=25.10
10/17/2025 5:32 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
10/17/2025 5:32 AM UTC
sha256:637559ec157d3f7a1f2986a05fb7a3d6048e44ece3fda9ec2cb222ef0ed8b697RUN
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
sha256:5a870de124badece7fa5b6d2c8587f4cb6d1f3382368d12f1da1b18b7c7fd5cbCOPY
pinned/25.10-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
10/17/2025 5:31 AM UTC
sha256:a3963738d44e6863b2e4da69182ee565ebdf72971060d75bdd8c3279f91f07fbRUN
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
sha256:0f96b2450ea37861a08c3198011bf6f6665125b82cbb1cd0fa3c058091e4fd16COPY
/opt/transformer-engine /opt/transformer-engine
10/17/2025 5:30 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
10/17/2025 5:30 AM UTC
sha256:5f97eef9460fb608f4e580d8ed26237656335fcd2f9b80075cb2c328cae90f33RUN
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
sha256:5b99269a23596d809e67d6da65238f31a894011450a9dd544ef99e7f003e04c1RUN
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
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1RUN
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
...

NVIDIA uses cookies to improve your experience on our web site. We and our third-party partners also use cookies and other tools to collect and record information you provide as well as information about your interactions with our websites for performance improvement, analytics, and to assist in marketing efforts. By clicking "Accept All", you consent to our use of cookies and other tools as described in our Cookie Policy. You can manage your cookie settings by clicking on "Manage Settings." By continuing to use this site or by clicking one of the buttons below, you agree to our Terms of Service (which contains important waivers). Please see our Privacy Policy for more information on our privacy practices.