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
b162e15a48e2171973d62f05fbe1dc76964bfedcaf29b113f85b718704dfa095CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
01/22/2025 8:37 AM UTC
2a8318708c2ea8e914c4dc2bbe7579a0d2609df2fd90378e764306ad43018256ENV
NVIDIA_PRODUCT_NAME=JAX
01/22/2025 8:37 AM UTC
dd0945b73969293a86bb3d702cf6484286508f61d7e222ee03001f416ddf3984RUN
URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 BUILD_PATH_JAXLIB=/opt/jaxlibs pip install "ml_dtypes>=0.5.0"
01/22/2025 8:37 AM UTC
a802608c8a0611f9d90efa4534f10bcec04c3c6add347d34666c10838729f4dcRUN
URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
01/22/2025 8:37 AM UTC
dff3a91c95581937346561ad4bebd70145d777931aa466694762d57124261b6fRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 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
01/22/2025 8:35 AM UTC
6e309762a44ef7b7103b4dfd67ab950afc5fbc3fd3bbe566a0887deae8b5136aCOPY
/opt/transformer-engine /opt/transformer-engine
01/22/2025 8:35 AM UTC
d1d9b2325ec48eee77d0c1558a9a3809605c993eaef74a53eee6bcecf711d111ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
01/22/2025 8:35 AM UTC
fbe342372320475c5663cc48cd78b45340b556065095d5adc0c90bf0664c8c99ENV
NVTE_FRAMEWORK=jax
01/22/2025 8:35 AM UTC
c76f72a3f9942cff4015936a21c2191a7fa5db85b7c22f5c1a414a17ca8f7e82RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit
01/22/2025 8:35 AM UTC
f811f46c6aab59f43942eb8c02de0ac66c6ddd7d7b50213a0f727e4607641cc9RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#53bde74afb8fc27270818226b5a27bd7cf195ac9 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-01-22 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}" >> /opt/pip-tools.d/requirements-jax.in for pkg in \
"optax==0.2.4" \
"orbax-checkpoint==0.10.2" \
"orbax-export==0.0.6" \
; do echo "$pkg" >> /opt/pip-tools.d/requirements-jax.in done EOF # buildkit
01/22/2025 8:35 AM UTC
...