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
871ca21b1136970482e7d81ef00bbff95a9b9035fecd7fd815c2e5f54bb15c49CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
04/16/2025 9:50 PM UTC
0fbafab6aa3ee63355aebc42f99b1cd58554d9d531f2babb77d63662b58e204fRUN
NVIDIA_BUILD_ID=159526777 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb SRC_PATH_MAXTEXT=/opt/maxtext pip install "werkzeug>=3.0.3"
04/16/2025 9:50 PM UTC
3188b9ba6bcf43a14302add85d45131f1e4f81775d69a64ceb8fa0fc092b2ce2RUN
NVIDIA_BUILD_ID=159526777 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
04/16/2025 9:50 PM UTC
760a108871ab085760358c838e605123c5ac6fca88ba8517b0f6dc15a9dac26eADD
test-maxtext.sh /usr/local/bin
04/16/2025 9:47 PM UTC
68c3d00c951c47eaf31b12622f17c80343509dfe09cdc443eb744897131b9cdcRUN
RUN |3 NVIDIA_BUILD_ID=159526777 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} echo "-r ${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in # specify some restrictions to speed up the build and # avoid pip to download and check all available versions of packages for pattern in \
"s|absl-py|absl-py>=2.1.0|g" \
"s|protobuf==3.20.3|protobuf>=3.19.0|g" \
"s|tensorflow-datasets|tensorflow-datasets>=4.8.0|g" \
"s|sentencepiece==0.1.97|sentencepiece>=0.2|g" \
"s|tensorflow-text>=2.13.0|tensorflow-text==2.18.1|g" \
"s|tensorflow>=2.13.0|tensorflow==2.18.1|g" \
; do sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt; done # add new line in case requirements.txt does not end with a new line echo >> ${SRC_PATH_MAXTEXT}/requirements.txt for requirement in \
"tensorflow-metadata>=1.15.0" \
"seqio@git+https://github.com/google/seqio.git" \
; do echo "${requirement}" >> ${SRC_PATH_MAXTEXT}/requirements.txt done EOF # buildkit
04/16/2025 9:47 PM UTC
79d61d033076a2f5c8abc88c4cec3bfe021365cd34316b99f056fa2b21b3e436ARG
SRC_PATH_MAXTEXT=/opt/maxtext
04/16/2025 9:47 PM UTC
0b024f85170510e5e1d651fc92dbae026cb1b868872344be61ac6c9f05f5a93dARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb
04/16/2025 9:47 PM UTC
e9a220dedef6fbeb54cc50c04fd7b414a201ff805725923c67aba6fcebca851aENV
NVIDIA_BUILD_ID=159526777
04/16/2025 9:47 PM UTC
115bc8e6eadaf450d9dc3f87cb556f7c6e8f41bcd9f45cf7c5df604725eff277ARG
NVIDIA_BUILD_ID=159526777
04/16/2025 9:47 PM UTC
6eded15089ce9a219505a68a72fa0d77bbcf4a4db81256a32b875b3a88da795bRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 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-04-16 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/16/2025 9:40 PM UTC
...