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
bb3b8fce8236b00c30f6c1c7060e877cef7a711f4c33722b6beb3474f2720023CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
04/23/2026 5:10 PM UTC
4a8b26fe9009d47fab7fa777980fdcefcd2426a4250774ecea689aeea6d20692ENV
NVIDIA_PRODUCT_NAME=JAX
04/23/2026 5:10 PM UTC
c44725a2662d09a02ac854ba999c3669030ca12cabfa07405d4a7a7c2f81eacbRUN
URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 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-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
04/23/2026 5:10 PM UTC
2b0f8d2e95238f63dc581b193094322eff5aff39324a02f0c713a967da19c62dCOPY
pinned/26.04-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
04/23/2026 5:09 PM UTC
3b7cf7473172309fffb81c5210afae958c790e0b56a6bc5ac0d581f39ad1176eRUN
URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 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-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
04/23/2026 5:08 PM UTC
dba96232fe0bcbaced2b4dd0f902f7727c779497723ee1d78e8d41dc1cda543fRUN
URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 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-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
04/23/2026 5:08 PM UTC
0368cb710b50d9872432c2f3402883760559240dc412823bc1c347c94fdc81a4ADD
install-nsys-jax.sh /usr/local/bin
04/23/2026 5:08 PM UTC
00e0d6615002eb4e8092fe4263202342441b6ea8aa7c3e14eeee3b30879b0e1dCOPY
nsys_jax/ /opt/nsys_jax
04/23/2026 5:08 PM UTC
8738009b10c3d60ba84a7e1f6fff77df446ba0ed53a029084c2151b4b829a42cENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
04/23/2026 5:08 PM UTC
a0bf8ba888c169b8ba64d0029c672f2d87e17c173f169563c7893fa96f4c13bdRUN
RUN |8 URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 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-04-23 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
04/23/2026 5:08 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.