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.06
06/05/2026 8:56 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
06/05/2026 8:56 PM UTC
sha256:edd0da588a35964286e03f54eb84ab2e465ca9fd43dd6b4d56025f624c9f5b1dRUN
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
sha256:ce97c6f1308dbfec26026121f0749ae26f96ddb06d7d272b72d5f1f9d049450dCOPY
pinned/26.06-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
06/05/2026 8:55 PM UTC
sha256:57e85b3d08ea54d40e88b96e309c8328b3dd92fce10926318187605097b5212dRUN
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
sha256:9f4d6f417c02ff07360066397d6d415b24963ebf2b40b755e14188007eb5c2cbRUN
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
sha256:162eab82830046414e4b46c68061febb6392684b340b892b6f55701987a318cfADD
install-nsys-jax.sh /usr/local/bin
06/05/2026 8:54 PM UTC
sha256:30574297615ea3e0ae6cbee5687154cfb14436dd10fe4c73a4e53f57cdcf7ad5COPY
nsys_jax/ /opt/nsys_jax
06/05/2026 8:54 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
06/05/2026 8:54 PM UTC
sha256:87ac31fb9ba86043ddefb2b379602ed2c24691b47ebb26812c19d6d7c784150cRUN
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.