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
04/16/2025 9:50 PM UTC
sha256:59f1ffe006963baf1c24e22d65c3900c7911cd5c7306fab7d7579dcdea80a822RUN
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
sha256:b1cb2010f80af6261709be63b779653d241967278150012a2e7e1833a5467db2RUN
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
sha256:9c30d93a32ca63f1474f040a04701a6aca2d5b4432d168ed9e56ae1979e66409ADD
test-maxtext.sh /usr/local/bin
04/16/2025 9:47 PM UTC
sha256:d144878496e4e06314f5677c054e3a96956fce96f783d5c413bf5a1d9e9886c8RUN
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
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
SRC_PATH_MAXTEXT=/opt/maxtext
04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#0b1fde4513d9b57c25dc3651b2f3c66497d800bb
04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=159526777
04/16/2025 9:47 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
NVIDIA_BUILD_ID=159526777
04/16/2025 9:47 PM UTC
sha256:b0c7621c7577f60feee0d061232b6c8cd6d382db48480c38449313ab0deda8e7RUN
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
...