Skip to main content
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
    09/16/2026 12:59 AM UTC
    sha256:0051fc262966ca54fc896aa6feb217aeffee262845c643507a6ca76c3efa67c8RUN
    NVIDIA_BUILD_ID=68066583 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0c794a551b264dd560f375cc3c053a2949ea3045 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c pip-finalize.sh
    09/16/2026 12:59 AM UTC
    sha256:0f4c5c008c67169a257d88e17e73e08db3de3bf5dc20f77584a9aa0a7ba3a6a4COPY
    pinned/26.09-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
    09/16/2026 12:57 AM UTC
    sha256:1b915858cecb41caa18920c37108ef347eb8917d6844fd8f84af1f2e4be22485ADD
    test-maxtext.sh /usr/local/bin
    09/16/2026 12:57 AM UTC
    sha256:53905f992e12e894119e33245d9fd08bad700196dbe3e77e55f75134a6143dd2RUN
    RUN |3 NVIDIA_BUILD_ID=68066583 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0c794a551b264dd560f375cc3c053a2949ea3045 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -exuo pipefail # Avoid old tensorboard constraint from google-cloud-aiplatform[tensorboard]. sed -i "/^cloud-accelerator-diagnostics/d" \
    ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/base_requirements/requirements.txt # remove custom build hook that conflicts with our dependencies introduced in # https://github.com/AI-Hypercomputer/maxtext/commit/de886f9a918d17cfb025a353e7da2bc4d64c26b5 sed -i -e '/^\[tool\.hatch\.build\.targets\.wheel\.hooks\.custom\]$/d' \
    -e '/^path = "build_hooks\.py"$/d' \
    ${SRC_PATH_MAXTEXT}/pyproject.toml EOF # buildkit
    09/16/2026 12:57 AM UTC
    sha256:0cd8c831306cb29b80cde2d6cf2102afe006d4e72af8e9f3854b4dc71fa159bcRUN
    RUN |3 NVIDIA_BUILD_ID=68066583 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0c794a551b264dd560f375cc3c053a2949ea3045 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -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/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in # MaxText's base_requirements/requirements.txt (above) does not include google-jetstream,
    # but inference code paths (maxtext.inference.maxengine, maxtext.input_pipeline.packing.prefill_packing)
    # import the `jetstream` package at module load time and DECOUPLE_GCLOUD=TRUE does not
    # unblock prefill_packing. Now re-add this dependency
    grep -E '^google-jetstream' ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/requirements.txt \
      >> /opt/pip-tools.d/requirements-maxtext.in
    EOF # buildkit
    09/16/2026 12:57 AM UTC
    sha256:3b8597adb8d8c201394f7ed02607a29606aec487f0d6a610cc7f2b550ecccff1RUN
    RUN |3 NVIDIA_BUILD_ID=68066583 URLREF_MAXTEXT=https://github.com/google/maxtext.git#0c794a551b264dd560f375cc3c053a2949ea3045 SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
    09/16/2026 12:57 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    SRC_PATH_MAXTEXT=/opt/maxtext
    09/16/2026 12:57 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    URLREF_MAXTEXT=https://github.com/google/maxtext.git#0c794a551b264dd560f375cc3c053a2949ea3045
    09/16/2026 12:57 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_BUILD_ID=68066583
    09/16/2026 12:57 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.