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:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_JAX_VERSION=25.01
    		01/22/2025 8:37 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_PRODUCT_NAME=JAX
    		01/22/2025 8:37 AM UTC
    sha256:9188dceb8bc7549b4d04613a435dda8d10c395822c9b8cae403c5ac49dca8910RUN
    URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 BUILD_PATH_JAXLIB=/opt/jaxlibs pip install "ml_dtypes>=0.5.0"
    		01/22/2025 8:37 AM UTC
    sha256:2965b920f9b6c2fefb1da8a7e0eda33cb8c1af59c19b615a823166d2270f4053RUN
    URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
    		01/22/2025 8:37 AM UTC
    sha256:18a3c8e1b70c39c2bc35d24c7f3d2269aa28f11a970532788184950530fb2ceaRUN
    RUN |7 URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 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
    		01/22/2025 8:35 AM UTC
    sha256:a1e2921e4e358da57593ae23ec9346b2769a31e9b12b6c9941f3af75ea859ccbCOPY
    /opt/transformer-engine /opt/transformer-engine
    		01/22/2025 8:35 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
    		01/22/2025 8:35 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVTE_FRAMEWORK=jax
    		01/22/2025 8:35 AM UTC
    sha256:66df50f7aa22bdb405b97cf2499b53bcd7071f82a931d9932acac55487534f79RUN
    RUN |7 URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 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
    		01/22/2025 8:35 AM UTC
    sha256:09dd72e2df6ff6c6d7c11cae3ac3d9d3dbc447c9e491fbfd1bbe4596785ba211RUN
    RUN |7 URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 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}" >> /opt/pip-tools.d/requirements-jax.in for pkg in \
"optax==0.2.4" \
"orbax-checkpoint==0.10.2" \
"orbax-export==0.0.6" \
; do echo "$pkg" >> /opt/pip-tools.d/requirements-jax.in done EOF # buildkit
    		01/22/2025 8:35 AM UTC
    ...