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
13a5f1b09abbb499774fb1a9235e9b299aba2ea74fd018d4722a1324db9941baCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
06/05/2026 8:56 PM UTC
8c0341b901e08d4443763ed239ac07ebf66603023b22030c864d6073e81a4fbeENV
NVIDIA_PRODUCT_NAME=JAX
06/05/2026 8:56 PM UTC
222dcb7aa8683d0a794795bd9c9f76c0407fde6c32e17706c9a4b7a32d0f5253RUN
URLREF_FLAX=https://github.com/google/flax.git#d044de84be4a6da55a0b36c478c429f5d6eb4fda 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-06-05 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
06/05/2026 8:56 PM UTC
2363626a3682237b7fdeef7e2305006519ca5270e7be759d34a8a6290f3dad74COPY
pinned/26.06-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
06/05/2026 8:55 PM UTC
e17d4997273a9b42455318b37b547540fa2931d6912f1517eb352e3dc1cd9328RUN
URLREF_FLAX=https://github.com/google/flax.git#d044de84be4a6da55a0b36c478c429f5d6eb4fda 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-06-05 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
06/05/2026 8:54 PM UTC
57a5b43d8e256ab44407e3763e14e3db217833314679a60a6d278364f9c7b9c1RUN
URLREF_FLAX=https://github.com/google/flax.git#d044de84be4a6da55a0b36c478c429f5d6eb4fda 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-06-05 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
06/05/2026 8:54 PM UTC
8106b92520acf1b4ce14b6315118cc85c2246681010028d3cbeccd1be453d310ADD
install-nsys-jax.sh /usr/local/bin
06/05/2026 8:54 PM UTC
001b74bb2a83c9fe5c78688d25d4e40f6067163e3e73d9d4cf82e46b32f48a42COPY
nsys_jax/ /opt/nsys_jax
06/05/2026 8:54 PM UTC
80e3a419e5cb0797477233242294227d49b57c918c25aa689ef8b2196bdb4b0cENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
06/05/2026 8:54 PM UTC
0140529639fc160ce2f1e27a6b5418476af7029f8f601acbd625a77cc3d8dedcRUN
URLREF_FLAX=https://github.com/google/flax.git#d044de84be4a6da55a0b36c478c429f5d6eb4fda 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-06-05 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "nvidia-cutlass-dsl[cu13]" >> /opt/pip-tools.d/requirements-cutlass-dsl.in
06/05/2026 8:54 PM 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.