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
    08/21/2026 6:29 PM UTC
    sha256:6ac7e202900ebb11eed6f84b255e93f953fea690add4ac38a15edc762676ca51RUN
    NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c pip-finalize.sh
    08/21/2026 6:29 PM UTC
    sha256:dd0375255cdd0d4a5a56d2a8af971d4f2aaaf7e9e47ec5bfc277bebb45450ce4COPY
    pinned/26.08-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
    08/21/2026 6:26 PM UTC
    sha256:ac0df7e7a23df864d2b4eb2100e2565cc3ad4c84646054e0365fedef5c5ec052ADD
    test-maxtext.sh /usr/local/bin
    08/21/2026 6:26 PM UTC
    sha256:f65589167301410224d8b2e280b89595d7f57d821ab8961b6f75549c8a9d2bd7RUN
    RUN |3 NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada 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
    08/21/2026 6:26 PM UTC
    sha256:085f48e1907060031ec226bff0fc2dfb96bf71c2965a01c3d6ff40a0479619c8RUN
    RUN |3 NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada 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
    08/21/2026 6:26 PM UTC
    sha256:f2461c8c6c6a179fc49d84abf9d51e4e62c93cc964994561cdcf6eb42db3128dRUN
    RUN |3 NVIDIA_BUILD_ID=63903073 URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada SRC_PATH_MAXTEXT=/opt/maxtext /bin/bash -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
    08/21/2026 6:26 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    SRC_PATH_MAXTEXT=/opt/maxtext
    08/21/2026 6:26 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    URLREF_MAXTEXT=https://github.com/google/maxtext.git#be69fa63e0d7d2fa0852b4cadaea2abd1ee2eada
    08/21/2026 6:26 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_BUILD_ID=63903073
    08/21/2026 6:26 PM UTC
    ...