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
21b47114ce084c22d2f85163a44df7b5f85015d2938c910ceddcc0f1fe607fb4CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
08/08/2025 6:53 PM UTC
8148292bab2f68bc8c760e5b75541463131bec12996cbb6363b7d86510607673ENV
NVIDIA_JAX_VERSION=25.08
08/08/2025 6:53 PM UTC
c62080bc9fa48cf4cdc2942d9f1415d64633934bf8f08798a5acb16f2878d28aENV
NVIDIA_PRODUCT_NAME=JAX
08/08/2025 6:53 PM UTC
4ba4674ce6305d878528ff58296e9124c838310d717beaa48cc05bd86547b601RUN
URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
08/08/2025 6:53 PM UTC
b29bd6bc1193dc2c33acff95334abd772a47ef68514f8e0ef971fd25848b4880RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 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
08/08/2025 6:48 PM UTC
2e9449a703ba893b3ec8c394fcbccabb871a433b8bd738a89d9ec69e069099ceCOPY
/opt/transformer-engine /opt/transformer-engine
08/08/2025 6:48 PM UTC
40b7eb50968a27a48d9df13386a3c22822de93e06ab9a6935e96c4ec60ec320aENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
08/08/2025 6:48 PM UTC
31d85d526516f0525f1480534bf4c897c8d6c327f2193aebf7d43d966d421ba4RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} sed -i 's/orbax-checkpoint/orbax-checkpoint==0.11.19/' ${SRC_PATH_FLAX}/pyproject.toml echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit
08/08/2025 6:48 PM UTC
24e1ba1edab766fb05b9ed16862c3766148f9fea56f98a21a186af3f1aaa1d9dRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex for component in $(ls ${BUILD_PATH_JAXLIB}); do echo "-e file://${BUILD_PATH_JAXLIB}/${component}" >> /opt/pip-tools.d/requirements-jax.in; done echo "-e file://${SRC_PATH_JAX}[k8s]" >> /opt/pip-tools.d/requirements-jax.in EOF # buildkit
08/08/2025 6:48 PM UTC
86c637937d43f9a691d8df34c22557b555f80d37bb9bab4394754855aa834d63RUN
URLREF_FLAX=https://github.com/google/flax.git#v0.10.6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-08-08 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d
08/08/2025 6:48 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.