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
681df89033a4c70d666b987cece831df062b1740ef8da931c965f848218a1166CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
03/24/2026 8:35 AM UTC
090f6a67a446f76fb887f055fb6f93c1abdaf92798c14010c5aa7a8c4f8a0ab1ENV
NVIDIA_PRODUCT_NAME=JAX
03/24/2026 8:35 AM UTC
ebdc58c592d191456ac9fde0ff94a47a7a187802029b25b47b0f98b3c26d408dRUN
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
6c1a608c30244c84515e330e19d13b9d5bbef5d7686b62496a1123143d38ae24COPY
pinned/26.03-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
03/24/2026 8:34 AM UTC
a56fdb46e395565266f0b2dd954167d8dc4d3999df309238a19b02ff1bc22f74RUN
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
4a8ed4139a7eaf26298db35bd487b3a1a8e53f72e7b6e85666307a6bbc557b62ADD
install-nsys-jax.sh /usr/local/bin
03/24/2026 8:34 AM UTC
ff38f41b29e9fbffaba21b0eb3b043a05ea41b7b766c4d4b4d8b59fb8b72c6d6COPY
nsys_jax/ /opt/nsys_jax
03/24/2026 8:34 AM UTC
4fc2b43ed85af95e9df9c7924dbab4e067b47275d40cf273e7c55c18565c6358ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
03/24/2026 8:34 AM UTC
605130ff4ac6ef5fcc92cfcf11ebb53bbb09230df5c241f1d0fe7766b44db90eRUN
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
8e5f5c274633746c901c7c3f627a2dcca1213f3f9a6e20c4ff6a4a0886745a33COPY
/opt/transformer-engine /opt/transformer-engine
03/24/2026 8:34 AM UTC
...