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
479d3fcf45ea4df12280a096d3663fe352d8e6bc430f6075f246da09d945202cCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
08/08/2025 7:03 PM UTC
57944da5b39573d52c4820b89564477a17eabe72bd0bc9a26734ad2f40ed50ebRUN
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
cdcc70674ba7a21acaf2cf3effce40a9d78214aa4525afa1e08f24178729ef40ADD
test-maxtext.sh /usr/local/bin
08/08/2025 6:59 PM UTC
c1cdebbffa53c94869cfbd000e9134d6230f24c07ad091c3b47b88b00b3bd71dRUN
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
5fb9a5713a1a68307baf90ffd117ef968d2b6c8ed9b55d63e53a52b74a17794bRUN
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
53e462cd5613aecd1fe2a8ada5504c18bec041f5387bee2138287ecdb403a93dRUN
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
06538e63ff5f5bd3e23d29caf6175716f775608c92048606d028049d18eebdbfRUN
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
5888d8eb514427dda319429b3b244d8735fc1ff922f6bde50f3c3277ae9812bfRUN
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
818cf9af4ec2e540ac67813e1688e58d98be3f67076cfeb83da1065bde48444fARG
SRC_PATH_JETSTREAM=/opt/jetstream
08/08/2025 6:59 PM UTC
3c8051de3c0cc8e24d0d3b1778f4f4f6283d71793178820fc6f578796b107eb7ARG
SRC_PATH_MAXTEXT=/opt/maxtext
08/08/2025 6:59 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.