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:ac3611e021618a3328963fb5cbe26c9b0e2a5ccb0ed0a948831cc762fdaeca22ADD
CHANGELOG.md /
10/13/2026 11:21 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_JAX_VERSION=26.10
10/13/2026 11:21 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
10/13/2026 11:21 PM UTC
sha256:dd638201aba1fdd695f1f105edb85cc0ae9502ac93f8ed4263e3decf943bcd30RUN
URLREF_FLAX=https://github.com/google/flax.git#4f42ed069731c83ef27fccb0cc42647684f80109 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-10-13 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c pip-finalize.sh
10/13/2026 11:21 PM UTC
sha256:9411b82dfba0ed917ebd5943988cf59650d9683281410833aa3c575f53ce76ccCOPY
pinned/26.10-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
10/13/2026 11:20 PM UTC
sha256:49d37131abf79d2bc0d90f7dda9896285e3b8c877cc149c3d0eb2dbb1931f7bfRUN
URLREF_FLAX=https://github.com/google/flax.git#4f42ed069731c83ef27fccb0cc42647684f80109 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-10-13 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
10/13/2026 11:19 PM UTC
sha256:c570fde6bd8096719b5c20d34f5165bcf22be233a740c31ba3ae2853cc5545d8RUN
URLREF_FLAX=https://github.com/google/flax.git#4f42ed069731c83ef27fccb0cc42647684f80109 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-10-13 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
10/13/2026 11:19 PM UTC
sha256:e0aa55cf87293e7420adf532e93df516484ac668259ea95bb347de9d2fc7b8caADD
install-nsys-jax.sh /usr/local/bin
10/13/2026 11:19 PM UTC
sha256:217a95a74fc457086f5c1b859038a08e435abdaab06c2641efae03fc1aa9b909COPY
nsys_jax/ /opt/nsys_jax
10/13/2026 11:19 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
10/13/2026 11:19 PM UTC
1-10 of 150 items

Modal Content

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.

Manage Settings