NVIDIA
JAX
Container
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
    ...

    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.