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
    02/26/2026 1:21 AM UTC
    sha256:6713f648e8097dfc22da1f654c533c50427a36e5a54b2f9ee11e6863ae4fb9bbRUN
    NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh
    02/26/2026 1:21 AM UTC
    sha256:54bd4758e2aea5512b7d6192376dacfc0fd1fcf64bc4e3987492e81754e04d07COPY
    pinned/26.02-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
    02/26/2026 1:16 AM UTC
    sha256:526474a8ff5a5892f3b454ec51a0d87ed1b86bab1328a745349580a094f545cfADD
    test-maxtext.sh /usr/local/bin
    02/26/2026 1:16 AM UTC
    sha256:a6b328f241ffa022e32187136a56aaecbe24ee17bd63f9428042e0656600b016RUN
    RUN |3 NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail for pattern in \
    "s|^tensorflow$|tensorflow==2.18.1|g" \
    "s|^tensorflow-text$|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}/dependencies/requirements/base_requirements/requirements.txt done # 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
    02/26/2026 1:16 AM UTC
    sha256:d241a2f602c6b04132e8d6b59aaa7369bd38c2f9f510ba231cd2baf181e96062RUN
    RUN |3 NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-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}/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
    02/26/2026 1:16 AM UTC
    sha256:47966881c48ee8335283c790dfa6b6ae49ae3df297908a0fce3643be9a7764e3RUN
    RUN |3 NVIDIA_BUILD_ID=44841511 URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit
    02/26/2026 1:16 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    SRC_PATH_MAXTEXT=/opt/maxtext
    02/26/2026 1:16 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    URLREF_MAXTEXT=https://gitlab-ci-token:glcbt-64_i6uBaS5CnRQfYi215sxF@gitlab-master.nvidia.com/dl/jax/maxtext.git#26.02-devel
    02/26/2026 1:16 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    NVIDIA_BUILD_ID=44841511
    02/26/2026 1:16 AM UTC
    ...