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=26.03
03/24/2026 8:35 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
03/24/2026 8:35 AM UTC
sha256:7aa2d1be4ed545cfce81d31c7dfc3a69bb8b467b6731b5415f8423381c1cc89eRUN
URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-03-24 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
03/24/2026 8:35 AM UTC
sha256:5a1ef08b58eebec6c02084b922f3a4eab4a9640877a05deca53d05fe75c0168aCOPY
pinned/26.03-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
03/24/2026 8:34 AM UTC
sha256:e1f6f68e86d5c38bc484e060c8af5e6844249e8b624096c7655dade08b821152RUN
URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-03-24 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
03/24/2026 8:34 AM UTC
sha256:eba7adb811cd4767fdabece617c0bf2b90af54276f53dad7b047c447971935eaADD
install-nsys-jax.sh /usr/local/bin
03/24/2026 8:34 AM UTC
sha256:7dbb99dfd933dd6fcb5d4ced227c22504721b9ffaa65b02a4c8019e796fd48e5COPY
nsys_jax/ /opt/nsys_jax
03/24/2026 8:34 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
03/24/2026 8:34 AM UTC
sha256:2613f10d65963ec0d6c2621f6b8cbb874139e5f93ea80f1e3372dfdd93b90319RUN
RUN |8 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 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-03-24 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
03/24/2026 8:34 AM UTC
sha256:aeff0e3652da8e4493da016710327cc0feeceafb5a95873a3c2ebe8ae9c5e578COPY
/opt/transformer-engine /opt/transformer-engine
03/24/2026 8:34 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.