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.05
05/21/2026 10:50 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
05/21/2026 10:50 AM UTC
sha256:932ab85622293720639c690caaa4a30460e6b200d3fa03260128f2d9a14550afRUN
URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
05/21/2026 10:50 AM UTC
sha256:e380d606fefb17da12f0fae4ae293687d6c49483b9fb84bed0c58930737f9472COPY
pinned/26.05-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
05/21/2026 10:49 AM UTC
sha256:0a88b5dbfda400b4505fc266d771ecfc0f71df1878426697ed5089b8a05291cdRUN
URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
05/21/2026 10:48 AM UTC
sha256:813facc106ea9c197ffa61ec424e8b8307547d3ac5912381d1bff07247b295feRUN
URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
05/21/2026 10:48 AM UTC
sha256:827b1ac4dca9e6baf6b7fc7cc7b151cb9779a6f2eb3912ba91ec48847dfeed55ADD
install-nsys-jax.sh /usr/local/bin
05/21/2026 10:48 AM UTC
sha256:341e33fd1ab0295c1f7e9550213203bda9118c70a3587358144cd993aafdf7a4COPY
nsys_jax/ /opt/nsys_jax
05/21/2026 10:48 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
05/21/2026 10:48 AM UTC
sha256:001bb1d3a63507d7e04a97a40bf52ed995b7795c3c42296a6ae992f27c99d5f1RUN
URLREF_FLAX=https://github.com/google/flax.git#901efc34b5d34a9c379c08dcbc3330fb1d7a042c 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-05-21 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "nvidia-cutlass-dsl[cu13]" >> /opt/pip-tools.d/requirements-cutlass-dsl.in
05/21/2026 10:48 AM UTC
...