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
224da9fe3b057ab00adf92a3c2ebb47ebec329409faa8e84e6dc31c421beca71CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
01/22/2025 8:57 AM UTC
db2e967b23cf5e51a518284309032539ae68cebc2368dbfa66f7142016efe82dRUN
NVIDIA_BUILD_ID=135630783 URLREF_MAXTEXT=https://github.com/google/maxtext.git#4651cb3c73de92e528a98281be29c05b27aaf501 URLREF_TFTEXT=https://github.com/tensorflow/text.git#master SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_TFTEXT=/opt/tensorflow-text URLREF_JETSTREAM=https://github.com/google/jetstream.git#9ca442137fec58bab819b7e561f0d204ebda4767 SRC_PATH_JETSTREAM=/opt/jetstream pip install "ml_dtypes>=0.5.0"
01/22/2025 8:57 AM UTC
24f028c537222fc63275fe430823d50242088614907eaefaa9534893b6c4d06eRUN
NVIDIA_BUILD_ID=135630783 URLREF_MAXTEXT=https://github.com/google/maxtext.git#4651cb3c73de92e528a98281be29c05b27aaf501 URLREF_TFTEXT=https://github.com/tensorflow/text.git#master SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_TFTEXT=/opt/tensorflow-text URLREF_JETSTREAM=https://github.com/google/jetstream.git#9ca442137fec58bab819b7e561f0d204ebda4767 SRC_PATH_JETSTREAM=/opt/jetstream pip install "werkzeug>=3.0.3"
01/22/2025 8:57 AM UTC
4397a115c98a577ab79e5a9b4ce5da94ffb6d818e383aa2c359fe6ed9f1286a8RUN
NVIDIA_BUILD_ID=135630783 URLREF_MAXTEXT=https://github.com/google/maxtext.git#4651cb3c73de92e528a98281be29c05b27aaf501 URLREF_TFTEXT=https://github.com/tensorflow/text.git#master SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_TFTEXT=/opt/tensorflow-text URLREF_JETSTREAM=https://github.com/google/jetstream.git#9ca442137fec58bab819b7e561f0d204ebda4767 SRC_PATH_JETSTREAM=/opt/jetstream pip-finalize.sh
01/22/2025 8:57 AM UTC
41614e462f44cf2dbe0f28cf1d70d09ca33f4de07e0347309edbc0bca9aee753RUN
NVIDIA_BUILD_ID=135630783 URLREF_MAXTEXT=https://github.com/google/maxtext.git#4651cb3c73de92e528a98281be29c05b27aaf501 URLREF_TFTEXT=https://github.com/tensorflow/text.git#master SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_TFTEXT=/opt/tensorflow-text URLREF_JETSTREAM=https://github.com/google/jetstream.git#9ca442137fec58bab819b7e561f0d204ebda4767 SRC_PATH_JETSTREAM=/opt/jetstream sed -i "s/ml_dtypes>=0.5.0/ml_dtypes>=0.4.0/" /opt/jax/build/requirements.in &&
  sed -i "s/ml_dtypes>=0.5.0/ml_dtypes>=0.4.0/" /opt/jax/setup.py
01/22/2025 8:53 AM UTC
f78f156802aea0ffab864b47ba5f466ef19f887895f877d79b854c94fa48ef7fRUN
RUN |7 NVIDIA_BUILD_ID=135630783 URLREF_MAXTEXT=https://github.com/google/maxtext.git#4651cb3c73de92e528a98281be29c05b27aaf501 URLREF_TFTEXT=https://github.com/tensorflow/text.git#master SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_TFTEXT=/opt/tensorflow-text URLREF_JETSTREAM=https://github.com/google/jetstream.git#9ca442137fec58bab819b7e561f0d204ebda4767 SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} echo "-e file://${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-jetstream.in EOF # buildkit
01/22/2025 8:53 AM UTC
4ea7f73e56af9835da06b92e4b1e43188ce2086dfda9b206dc7186e9401cd9f9ARG
SRC_PATH_JETSTREAM=/opt/jetstream
01/22/2025 8:53 AM UTC
1108c170adad3592599de74f46d4125a96a37ca4336c2dc8f94f99dcd0062d21ARG
URLREF_JETSTREAM=https://github.com/google/jetstream.git#9ca442137fec58bab819b7e561f0d204ebda4767
01/22/2025 8:53 AM UTC
8e446c966fbca61aacfdf9157e1cc6d5aa1c49358fa56470fdaff9be33349563ADD
test-maxtext.sh /usr/local/bin
01/22/2025 8:53 AM UTC
1599e48d09b29b41337d2f77c22f86eccfc2e582fd055d3b4ecf798cb8787b7aRUN
RUN |5 NVIDIA_BUILD_ID=135630783 URLREF_MAXTEXT=https://github.com/google/maxtext.git#4651cb3c73de92e528a98281be29c05b27aaf501 URLREF_TFTEXT=https://github.com/tensorflow/text.git#master SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_TFTEXT=/opt/tensorflow-text /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} sed -i '/google-jetstream/d' ${SRC_PATH_MAXTEXT}/requirements.txt 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|grain-nightly|grain|g" \
"s|tensorboardx|tensorboardx>=2.6.2|g" \
"s|sentencepiece==0.1.97|sentencepiece>=0.2|g" \
; do sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt; done 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
01/22/2025 8:53 AM 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.