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
    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.

  • 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
    ...