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:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
/opt/maxtext
08/08/2025 7:03 PM UTC
sha256:4c5f7677047a3aafdb07927863b863ba7f02b3c849849e74a6703d524acba6c7RUN
NVIDIA_BUILD_ID=33017886 URLREF_MAXTEXT=https://github.com/google/maxtext.git#3662540ee852d0d8f8333a36c04ddc0f1316ebfb URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8e73820993f77cfc8efe34eb2a73f5de98 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream pip-finalize.sh
08/08/2025 7:03 PM UTC
sha256:1f8d9a256c742d5af6e39e8d5552834d42206b84e81d61724f533094dc6ea57eADD
test-maxtext.sh /usr/local/bin
08/08/2025 6:59 PM UTC
sha256:8f3c967f44b45d22f0b44dd0f3d8e07750bb2d7d593af9e2a041bc747a625e6cRUN
RUN |5 NVIDIA_BUILD_ID=33017886 URLREF_MAXTEXT=https://github.com/google/maxtext.git#3662540ee852d0d8f8333a36c04ddc0f1316ebfb URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8e73820993f77cfc8efe34eb2a73f5de98 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail echo >> ${SRC_PATH_MAXTEXT}/requirements.txt # add new line for requirement in \
"tensorflow-metadata>=1.15.0" \
"seqio@git+https://github.com/google/seqio.git" \
"werkzeug>=3.0.3" \
; do echo "${requirement}" >> ${SRC_PATH_MAXTEXT}/requirements.txt done EOF # buildkit
08/08/2025 6:59 PM UTC
sha256:38ab449821481d2da2361f45a7e27e778a6c3ceef544cbb803d554e3b0c08a4fRUN
RUN |5 NVIDIA_BUILD_ID=33017886 URLREF_MAXTEXT=https://github.com/google/maxtext.git#3662540ee852d0d8f8333a36c04ddc0f1316ebfb URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8e73820993f77cfc8efe34eb2a73f5de98 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail 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>=2.13.0|tensorflow==2.18.1|g" \
"s|google-cloud-aiplatform==1.61.0|google-cloud-aiplatform>=1.90.0|g" \
"s|orbax-checkpoint>=0.5.12|orbax-checkpoint==0.11.19|g" \
; do # tensorflow-cpu==2.19.0 is incompatible with tensorflow-text sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt done EOF # buildkit
08/08/2025 6:59 PM UTC
sha256:6cece2599ce4df7e3a72de275adf0c925b47f5d6a2e687e86d5804a8b24a6a40RUN
RUN |5 NVIDIA_BUILD_ID=33017886 URLREF_MAXTEXT=https://github.com/google/maxtext.git#3662540ee852d0d8f8333a36c04ddc0f1316ebfb URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8e73820993f77cfc8efe34eb2a73f5de98 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail sed -i '/^google-jetstream/d' ${SRC_PATH_MAXTEXT}/requirements.txt EOF # buildkit
08/08/2025 6:59 PM UTC
sha256:5ced8388f296d8d46b8517bb7a12f04ca69378fc3be23ceddb2ab2547cc3c479RUN
RUN |5 NVIDIA_BUILD_ID=33017886 URLREF_MAXTEXT=https://github.com/google/maxtext.git#3662540ee852d0d8f8333a36c04ddc0f1316ebfb URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8e73820993f77cfc8efe34eb2a73f5de98 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail echo "-r ${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in echo "-e file://${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
08/08/2025 6:59 PM UTC
sha256:f9be6fc255ec2771292f27eae75fd08e499d11332aa2bd000f118c5829dcbed8RUN
RUN |5 NVIDIA_BUILD_ID=33017886 URLREF_MAXTEXT=https://github.com/google/maxtext.git#3662540ee852d0d8f8333a36c04ddc0f1316ebfb URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8e73820993f77cfc8efe34eb2a73f5de98 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} EOF # buildkit
08/08/2025 6:59 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_JETSTREAM=/opt/jetstream
08/08/2025 6:59 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
08/08/2025 6:59 PM UTC
...