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:3bb2f5d1304ef81d2a2136eda2a5b4f67adf378103d439b1d220752404b65810ADD
CHANGELOG.md /
08/21/2026 6:21 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_JAX_VERSION=26.08
08/21/2026 6:21 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
08/21/2026 6:21 PM UTC
sha256:e52952a8c07b1bc4aea7755dfc2226e9180625544b09341ced4384b1597bc56dRUN
URLREF_FLAX=https://github.com/google/flax.git#a31129e2f8a2f5ca6ccc8492e536d2b28359f747 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-08-21 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c pip-finalize.sh
08/21/2026 6:21 PM UTC
sha256:873e203cfdd6bfec64e94b1a60820105859098197ac4c5e52a98fb215af3324cCOPY
pinned/26.08-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
08/21/2026 6:21 PM UTC
sha256:d101ceba42abaf1b13a8995249668b7d1aa06a7d453efd5b23d86145b8135f0fRUN
URLREF_FLAX=https://github.com/google/flax.git#a31129e2f8a2f5ca6ccc8492e536d2b28359f747 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-08-21 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
08/21/2026 6:20 PM UTC
sha256:3eef45762fd1e9aef0e5232a42353c046e2865b1f6b00459cf2b6ae8a1dea9f5RUN
URLREF_FLAX=https://github.com/google/flax.git#a31129e2f8a2f5ca6ccc8492e536d2b28359f747 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-08-21 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/bash -c install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
08/21/2026 6:20 PM UTC
sha256:22fbd881f10cabf573ffc358bcc88e6227b5073610045fdbbf747b44d0707fcbADD
install-nsys-jax.sh /usr/local/bin
08/21/2026 6:20 PM UTC
sha256:740348b1171fce7d02cd4567db4d564e4a24a4a50f84af6858d5632064dba5f5COPY
nsys_jax/ /opt/nsys_jax
08/21/2026 6:20 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
08/21/2026 6:20 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