Skip to main content
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:d3d4122b2bec6c98ef6c975dd37324b309345fb6e7b6c34fbb7fd25ac8570038ADD
    CHANGELOG.md /
    		09/16/2026 12:08 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_JAX_VERSION=26.09
    		09/16/2026 12:08 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_PRODUCT_NAME=JAX
    		09/16/2026 12:08 AM UTC
    sha256:e15f28d0454a1c97897dbe2814d10642fe3a5fd168d15ffc2dc460d3aa57e6d3RUN
    URLREF_FLAX=https://github.com/google/flax.git#c5f2e5b0f3bee8bbc398cf74cc036d8eaacdf9a9 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-09-15 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c pip-finalize.sh
    		09/16/2026 12:08 AM UTC
    sha256:1ca8ece19cd66e8bf8e309329c3ce8e091452663c4b4e072073fde7c69421d86COPY
    pinned/26.09-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
    		09/16/2026 12:07 AM UTC
    sha256:ce30de20ee461ac0bf666c28dc7a3b241a10fbf6d6abd28d102f520829befa03RUN
    URLREF_FLAX=https://github.com/google/flax.git#c5f2e5b0f3bee8bbc398cf74cc036d8eaacdf9a9 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-09-15 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
    		09/16/2026 12:06 AM UTC
    sha256:72414103e09bfbe8c4d66f4fbd6f49264b2cbcbc131ed00631f7d32ec57e1eedRUN
    URLREF_FLAX=https://github.com/google/flax.git#c5f2e5b0f3bee8bbc398cf74cc036d8eaacdf9a9 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-09-15 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
    		09/16/2026 12:06 AM UTC
    sha256:2c3bbe018ce82ea639dc647c3796e0030acd0592debf3287ee9b3010f3875025ADD
    install-nsys-jax.sh /usr/local/bin
    		09/16/2026 12:06 AM UTC
    sha256:727ad2045c011c4bcc599983bf3308473d18d2d5a441485202bd6054e8f1856aCOPY
    nsys_jax/ /opt/nsys_jax
    		09/16/2026 12:06 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    SRC_PATH_NSYS_JAX=/opt/nsys_jax
    		09/16/2026 12:06 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.