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
    10/17/2025 5:38 AM UTC
    sha256:e15213cd0e1e19ef34cc335c1cca0c28eca368d4a2cd27285a5d8563f058fae3RUN
    NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream pip-finalize.sh
    10/17/2025 5:38 AM UTC
    sha256:1fb119e6c84edca9ccf13ec9a0f9527ccb5f9aaeb299494860d5a36aa2717c0fCOPY
    pinned/25.10-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
    10/17/2025 5:36 AM UTC
    sha256:dc1502f5d478abf24dfbd4c7148f18287fb66699a4b0b3a976e1ef68e5558465ADD
    test-maxtext.sh /usr/local/bin
    10/17/2025 5:36 AM UTC
    sha256:556dbdc9c3323ea9ad71bb3c9a967b3306901e4b21c4fa6fcb5f4ea62069cb4fRUN
    RUN |5 NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail for pattern in \
    "s|tensorflow>=2.19.1|tensorflow==2.18.1|g" \
    "s|tensorboard>=2.19.0|tensorboard>=2.18,<2.19|g" \
    "s|tensorflow-text>=2.19.0|tensorflow-text==2.18.1|g" \
    ; do # tensorflow-cpu,tensorboard,tensorflow-text>=2.19.0 is incompatible with tensorflow==2.18.1 sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/pyproject.toml done EOF # buildkit
    10/17/2025 5:36 AM UTC
    sha256:8008f44ff29f4e2e813ae5509ae1bfaf498ff718b39c4cf00c546ac32a663ff2RUN
    RUN |5 NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 SRC_PATH_MAXTEXT=/opt/maxtext SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex -o pipefail echo "-e file://${SRC_PATH_MAXTEXT}" >> /opt/pip-tools.d/requirements-maxtext.in echo "-r ${SRC_PATH_MAXTEXT}/src/install_maxtext_extra_deps/extra_deps_from_github.txt" >> /opt/pip-tools.d/requirements-maxtext.in #echo "-e file://${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
    10/17/2025 5:36 AM UTC
    sha256:b1e826a886a9a6938bab91a61f92461c2b527c46636be2e39ec1b769268ec508RUN
    RUN |5 NVIDIA_BUILD_ID=36776647 URLREF_MAXTEXT=https://github.com/google/maxtext.git#8fdac10 URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8 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
    10/17/2025 5:36 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    SRC_PATH_JETSTREAM=/opt/jetstream
    10/17/2025 5:36 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    SRC_PATH_MAXTEXT=/opt/maxtext
    10/17/2025 5:36 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    URLREF_JETSTREAM=https://github.com/google/jetstream.git#29329e8
    10/17/2025 5:36 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.