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
0cd821663484e514c6145d3fd36008466ee433dafab7eb15179aaecc80685eecCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
05/21/2026 10:50 AM UTC
9711c9e303eda916b124e3f051df283d2aec8a77169772e28d12342b6f401295ENV
NVIDIA_PRODUCT_NAME=JAX
05/21/2026 10:50 AM UTC
bfa3b75177a66eac0523e877b6bb0c958b7663be3e1fa36678bbbfe951468190RUN
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
9a53c28ce2e6a6fb2805b00a474c7870523cdd6c9b37a34253dd46dc8371bb14COPY
pinned/26.05-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
05/21/2026 10:49 AM UTC
fa28844f6b8a66395879743c1484ffdc6bd81048efd03d79ca588cd278c3471aRUN
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
3e5e0f0d090e6ec7f9e2e368af6e4960dba694fc72a65d77a51c5c5c30363dbdRUN
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
05c228da1d2e042acfa2de3ed1d81c1b33e3c6fc975280b5a717b2c8e9f2f24aADD
install-nsys-jax.sh /usr/local/bin
05/21/2026 10:48 AM UTC
6277b69337d8f28792a93819f10bf4ca93615b885a2b5465b706456f485ffda9COPY
nsys_jax/ /opt/nsys_jax
05/21/2026 10:48 AM UTC
06a58bc668b16832871b7254e4f0a872eb5296c27fe11da28d8eae66d2ec0508ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
05/21/2026 10:48 AM UTC
a3905b73e376c37b63f91c8d287a98c9568f8e91700701d2dcfe16128b0812faRUN
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
...

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.