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
    04/23/2026 5:15 PM UTC
    sha256:085dd8d9afc01334f534f697f1e943f8a3301376d31eb39e3b18b0c1a4c87391RUN
    NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
    04/23/2026 5:15 PM UTC
    sha256:41e5181ca78c46738360314df044da6c0da28479e9ea04f6990d7d7244ba3aa5COPY
    pinned/26.04-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
    04/23/2026 5:14 PM UTC
    sha256:9a91ef40d740741a9597558061180e0177c3eda3b1ee1c59482685bfe15bc1f3ADD
    test-maxtext.sh /usr/local/bin
    04/23/2026 5:14 PM UTC
    sha256:9138061b14b96f79b13027e4167bd811c321fe1f859f4c9e7d2b88f4a427f23dRUN
    RUN |3 NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -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
    04/23/2026 5:14 PM UTC
    sha256:a6c1412f14ee3cc1d533f6b677c2b8bea5b0500c8037c01cdb0829bf1d8dc72fRUN
    RUN |3 NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext /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/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
    04/23/2026 5:14 PM UTC
    sha256:f19299bc21503897b83052481e9f7205f0135f42cd0437cae2daf7d37cd8adadRUN
    RUN |3 NVIDIA_BUILD_ID=49311194 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
    04/23/2026 5:14 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    SRC_PATH_MAXTEXT=/opt/maxtext
    04/23/2026 5:14 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_gQx25kMzKfwmNyLx5oSk@gitlab-master.nvidia.com/dl/jax/maxtext#26.04-devel
    04/23/2026 5:14 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_BUILD_ID=49311194
    04/23/2026 5:14 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.