chore: import upstream snapshot with attribution

ci/official/containers/ml_build/Dockerfile

@@ -0,0 +1,125 @@
+# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+################################################################################
+ARG BASE_IMAGE=ubuntu:22.04@sha256:445586e41c1de7dfda82d2637f5ff688deea9eb5f5812f8c145afacc35b9f0db
+FROM $BASE_IMAGE AS devel
+# See https://docs.docker.com/reference/dockerfile/#understand-how-arg-and-from-interact
+# on why we cannot reference BASE_IMAGE again unless we declare it again.
+################################################################################
+
+# Install devtoolset build dependencies
+COPY setup.sources.sh /setup.sources.sh
+COPY setup.packages.sh /setup.packages.sh
+COPY builder.packages.txt /builder.packages.txt
+
+RUN /setup.sources.sh && /setup.packages.sh /builder.packages.txt
+
+# Setup Python
+COPY setup.python.sh /setup.python.sh
+COPY builder.requirements.txt /builder.requirements.txt
+RUN /setup.python.sh python3.10 /builder.requirements.txt
+RUN /setup.python.sh python3.11 /builder.requirements.txt
+RUN /setup.python.sh python3.13 /builder.requirements.txt
+RUN /setup.python.sh python3.13-nogil /builder.requirements.txt
+RUN /setup.python.sh python3.14 /builder.requirements.txt
+RUN /setup.python.sh python3.14-nogil /builder.requirements.txt
+
+# Since we are using python3.12 as the default python version, we need to
+# install python3.12 last for now.
+# TODO(b/376338367): switch to pyenv.
+RUN /setup.python.sh python3.12 /builder.requirements.txt
+
+ARG INSTALL_NVIDIA_PACKAGES=false
+ARG NVIDIA_PACKAGES_FILE=cuda12.1_cudnn9.1.packages.txt
+COPY ${NVIDIA_PACKAGES_FILE} /nvidia.packages.txt
+COPY setup.sources.cudnn.sh /setup.sources.cudnn.sh
+
+ARG INSTALL_CMAKE_TOOLCHAIN=false
+ARG CMAKE_TOOLCHAIN_PACKAGES_FILE=cmake_toolchain.packages.txt
+COPY ${CMAKE_TOOLCHAIN_PACKAGES_FILE} /cmake_toolchain.packages.txt
+COPY setup.toolchain.cmake.sh /setup.toolchain.cmake.sh
+
+RUN if [ "${INSTALL_NVIDIA_PACKAGES}" = "true" ]; then \
+      echo "Installing Nvidia packages"; \
+      /setup.sources.cudnn.sh && /setup.packages.sh /nvidia.packages.txt; \
+    else \
+      echo "Nvidia packages will not be installed"; \
+    fi
+
+RUN if [ "${INSTALL_CMAKE_TOOLCHAIN}" = "true" ]; then \
+      echo "Installing CMake toolchain"; \
+      /setup.toolchain.cmake.sh && /setup.packages.sh /cmake_toolchain.packages.txt; \
+    else \
+      echo "CMake toolchain will not be installed"; \
+    fi
+
+# Setup links for TensorFlow to compile.
+# Referenced in devel.usertools/*.bazelrc.
+# Set python3.12 as the default python version.
+# TF does not support python3.13.
+RUN ln -sf /usr/bin/python3.12 /usr/bin/python3
+RUN ln -sf /usr/bin/python3.12 /usr/bin/python
+RUN ln -sf /usr/lib/python3.12 /usr/lib/tf_python
+
+# Link the compat driver to the location if available.
+RUN if [ -e "/usr/local/cuda/compat/libcuda.so.1" ]; then ln -s /usr/local/cuda/compat/libcuda.so.1 /usr/lib/x86_64-linux-gnu/libcuda.so.1; fi
+
+# Install various tools.
+# - bats: bash unit testing framework
+# - bazelisk: always use the correct bazel version
+# - buildifier: clean bazel build deps
+# - buildozer: clean bazel build deps
+# - gcloud SDK: communicate with Google Cloud Platform (GCP) for RBE, CI
+# - patchelf: Utility tool to modify existing ELF executables and libraries
+RUN git clone --branch v1.13.0 https://github.com/bats-core/bats-core.git && bats-core/install.sh /usr/local && rm -rf bats-core
+RUN wget https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64 -O /usr/local/bin/bazel && \
+    echo "e1508323f347ad1465a887bc5d2bfb91cffc232d11e8e997b623227c6b32fb76  /usr/local/bin/bazel" | sha256sum --check && \
+    chmod +x /usr/local/bin/bazel
+RUN wget https://github.com/bazelbuild/buildtools/releases/download/v8.2.1/buildifier-linux-amd64 -O /usr/local/bin/buildifier && \
+    echo "6ceb7b0ab7cf66fceccc56a027d21d9cc557a7f34af37d2101edb56b92fcfa1a  /usr/local/bin/buildifier" | sha256sum --check && \
+    chmod +x /usr/local/bin/buildifier
+RUN wget https://github.com/bazelbuild/buildtools/releases/download/v8.2.1/buildozer-linux-amd64 -O /usr/local/bin/buildozer && \
+    echo "04454a6a89c64c603027cc3371eb1c36e48727e04558e077c20ec37c9c2f831a  /usr/local/bin/buildozer" | sha256sum --check && \
+    chmod +x /usr/local/bin/buildozer
+
+RUN wget https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-483.0.0-linux-x86_64.tar.gz && \
+    echo "dd3ebdbd13bffa1997232ab5a68af7797594079b062fd41474c666b84165cc35  google-cloud-cli-483.0.0-linux-x86_64.tar.gz" | sha256sum --check && \
+    tar zxf google-cloud-cli-483.0.0-linux-x86_64.tar.gz && \
+    rm google-cloud-cli-483.0.0-linux-x86_64.tar.gz && \
+    google-cloud-sdk/install.sh --quiet && \
+    ln -s /google-cloud-sdk/bin/gcloud /usr/bin/gcloud
+ENV PATH="$PATH:/google-cloud-sdk/bin/:/usr/local/cuda/bin/"
+
+# Download and install gh CLI v2.49.0 from GitHub.
+RUN wget https://github.com/cli/cli/releases/download/v2.49.0/gh_2.49.0_linux_amd64.tar.gz && \
+    echo "3f7ef10a8ae6164220db20cf6d4fb96528c61d614f0803ccf39b1cabb6ae7263  gh_2.49.0_linux_amd64.tar.gz" | sha256sum --check && \
+    tar zxf gh_2.49.0_linux_amd64.tar.gz && \
+    mv gh_2.49.0_linux_amd64/bin/gh /usr/local/bin/gh && \
+    rm -rf gh_2.49.0_linux_amd64*
+
+# Download and install patchelf v0.18.0 from GitHub. The default Ubuntu focal
+# packages only provide the "0.10-2build1" version. We use patchelf to manipulate
+# certain shared libraries during the wheel building process (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/pip_package/build_pip_package.sh#L255-L262).
+# When we use Patchelf versions <0.12, those shared libraries end up with a
+# corrupted PT_NOTE program header. This was fixed in v0.12, see https://github.com/NixOS/patchelf/commit/43a33482b501b0f5ee9da312aabfca3806570cc9.
+RUN wget https://github.com/NixOS/patchelf/releases/download/0.18.0/patchelf-0.18.0-x86_64.tar.gz && \
+    echo "ce84f2447fb7a8679e58bc54a20dc2b01b37b5802e12c57eece772a6f14bf3f0  patchelf-0.18.0-x86_64.tar.gz" | sha256sum --check && \
+    tar -zxvf patchelf-0.18.0-x86_64.tar.gz -C /usr && \
+    rm -rf patchelf-0.18.0-x86_64.tar.gz
+
+# Don't use the bazel cache when a new docker image is created.
+RUN echo build --action_env=DOCKER_CACHEBUSTER=$(date +%s%N)$RANDOM >> /etc/bazel.bazelrc
+RUN echo build --host_action_env=DOCKER_HOST_CACHEBUSTER=$(date +%s%N)$RANDOM >> /etc/bazel.bazelrc

ci/official/containers/ml_build/README.md

@@ -0,0 +1,8 @@
+WIP ML Build Docker container for ML repositories (Tensorflow, JAX and XLA).
+
+This container branches off from
+/tensorflow/tools/tf_sig_build_dockerfiles/. However, since
+hermetic CUDA and hermetic Python is now available for Tensorflow, a lot of the
+requirements installed on the original container can be removed to reduce the
+footprint of the container and make it more reusable across different ML
+repositories.

ci/official/containers/ml_build/builder.devtoolset/build_devtoolset.sh

@@ -0,0 +1,198 @@
+#!/bin/bash -eu
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Builds a devtoolset cross-compiler targeting manylinux 2010 (glibc 2.12 /
+# libstdc++ 4.4) or manylinux2014 (glibc 2.17 / libstdc++ 4.8).
+
+VERSION="$1"
+TARGET="$2"
+
+case "${VERSION}" in
+devtoolset-7)
+  LIBSTDCXX_VERSION="6.0.24"
+  LIBSTDCXX_ABI="gcc4-compatible"
+  ;;
+devtoolset-9)
+  LIBSTDCXX_VERSION="6.0.28"
+  LIBSTDCXX_ABI="new"
+  ;;
+*)
+  echo "Usage: $0 {devtoolset-7|devtoolset-9} <target-directory>"
+  echo "Use 'devtoolset-7' to build a manylinux2010 compatible toolchain or 'devtoolset-9' to build a manylinux2014 compatible toolchain"
+  exit 1
+  ;;
+esac
+
+mkdir -p "${TARGET}"
+
+# Download glibc's shared and development libraries based on the value of the
+# `VERSION` parameter.
+# Note: 'Templatizing' this and the other conditional branches would require
+# defining several variables (version, os, path) making it difficult to maintain
+# and extend for future modifications.
+case "${VERSION}" in
+devtoolset-7)
+  # Download binary glibc 2.12 shared library release.
+  wget "http://old-releases.ubuntu.com/ubuntu/pool/main/e/eglibc/libc6_2.12.1-0ubuntu6_amd64.deb" && \
+      unar "libc6_2.12.1-0ubuntu6_amd64.deb" && \
+      tar -C "${TARGET}" -xvzf "libc6_2.12.1-0ubuntu6_amd64/data.tar.gz" && \
+      rm -rf "libc6_2.12.1-0ubuntu6_amd64.deb" "libc6_2.12.1-0ubuntu6_amd64"
+  # Download binary glibc 2.12 development library release.
+  wget "http://old-releases.ubuntu.com/ubuntu/pool/main/e/eglibc/libc6-dev_2.12.1-0ubuntu6_amd64.deb" && \
+      unar "libc6-dev_2.12.1-0ubuntu6_amd64.deb" && \
+      tar -C "${TARGET}" -xvzf "libc6-dev_2.12.1-0ubuntu6_amd64/data.tar.gz" && \
+      rm -rf "libc6-dev_2.12.1-0ubuntu6_amd64.deb" "libc6-dev_2.12.1-0ubuntu6_amd64"
+  ;;
+devtoolset-9)
+  # Download binary glibc 2.17 shared library release.
+  wget "http://old-releases.ubuntu.com/ubuntu/pool/main/e/eglibc/libc6_2.17-0ubuntu5.1_amd64.deb" && \
+      unar "libc6_2.17-0ubuntu5.1_amd64.deb" && \
+      tar -C "${TARGET}" -xvzf "libc6_2.17-0ubuntu5.1_amd64/data.tar.gz" && \
+      rm -rf "libc6_2.17-0ubuntu5.1_amd64.deb" "libc6_2.17-0ubuntu5.1_amd64"
+  # Download binary glibc 2.17 development library release.
+  wget "http://old-releases.ubuntu.com/ubuntu/pool/main/e/eglibc/libc6-dev_2.17-0ubuntu5.1_amd64.deb" && \
+      unar "libc6-dev_2.17-0ubuntu5.1_amd64.deb" && \
+      tar -C "${TARGET}" -xvzf "libc6-dev_2.17-0ubuntu5.1_amd64/data.tar.gz" && \
+      rm -rf "libc6-dev_2.17-0ubuntu5.1_amd64.deb" "libc6-dev_2.17-0ubuntu5.1_amd64"
+  ;;
+esac
+
+# Put the current kernel headers from ubuntu in place.
+ln -s "/usr/include/linux" "/${TARGET}/usr/include/linux"
+ln -s "/usr/include/asm-generic" "/${TARGET}/usr/include/asm-generic"
+ln -s "/usr/include/x86_64-linux-gnu/asm" "/${TARGET}/usr/include/asm"
+
+# Symlinks in the binary distribution are set up for installation in /usr, we
+# need to fix up all the links to stay within /${TARGET}.
+/fixlinks.sh "/${TARGET}"
+
+# Patch to allow non-glibc 2.12 compatible builds to work.
+sed -i '54i#define TCP_USER_TIMEOUT 18' "/${TARGET}/usr/include/netinet/tcp.h"
+
+# Download specific version of libstdc++ shared library based on the value of
+# the `VERSION` parameter
+case "${VERSION}" in
+devtoolset-7)
+  # Download binary libstdc++ 4.4 release we are going to link against.
+  # We only need the shared library, as we're going to develop against the
+  # libstdc++ provided by devtoolset.
+  wget "http://old-releases.ubuntu.com/ubuntu/pool/main/g/gcc-4.4/libstdc++6_4.4.3-4ubuntu5_amd64.deb" && \
+      unar "libstdc++6_4.4.3-4ubuntu5_amd64.deb" && \
+      tar -C "/${TARGET}" -xvzf "libstdc++6_4.4.3-4ubuntu5_amd64/data.tar.gz" "./usr/lib/libstdc++.so.6.0.13" && \
+      rm -rf "libstdc++6_4.4.3-4ubuntu5_amd64.deb" "libstdc++6_4.4.3-4ubuntu5_amd64"
+  ;;
+devtoolset-9)
+  # Download binary libstdc++ 4.8 shared library release
+  wget "http://old-releases.ubuntu.com/ubuntu/pool/main/g/gcc-4.8/libstdc++6_4.8.1-10ubuntu8_amd64.deb" && \
+      unar "libstdc++6_4.8.1-10ubuntu8_amd64.deb" && \
+      tar -C "/${TARGET}" -xvzf "libstdc++6_4.8.1-10ubuntu8_amd64/data.tar.gz" "./usr/lib/x86_64-linux-gnu/libstdc++.so.6.0.18"  && \
+      rm -rf "libstdc++6_4.8.1-10ubuntu8_amd64.deb" "libstdc++6_4.8.1-10ubuntu8_amd64"
+  ;;
+esac
+
+mkdir -p "${TARGET}-src"
+cd "${TARGET}-src"
+
+# Build a devtoolset cross-compiler based on our glibc 2.12/glibc 2.17 sysroot setup.
+case "${VERSION}" in
+devtoolset-7)
+  wget "http://vault.centos.org/centos/6/sclo/Source/rh/devtoolset-7/devtoolset-7-gcc-7.3.1-5.15.el6.src.rpm"
+  rpm2cpio "devtoolset-7-gcc-7.3.1-5.15.el6.src.rpm" |cpio -idmv
+  tar -xvjf "gcc-7.3.1-20180303.tar.bz2" --strip 1
+  ;;
+devtoolset-9)
+  wget "https://vault.centos.org/centos/7/sclo/Source/rh/devtoolset-9-gcc-9.3.1-2.2.el7.src.rpm"
+  rpm2cpio "devtoolset-9-gcc-9.3.1-2.2.el7.src.rpm" |cpio -idmv
+  tar -xvf "gcc-9.3.1-20200408.tar.xz" --strip 1
+  ;;
+esac
+
+# Apply the devtoolset patches to gcc.
+/rpm-patch.sh "gcc.spec"
+
+./contrib/download_prerequisites
+
+mkdir -p "${TARGET}-build"
+cd "${TARGET}-build"
+
+"${TARGET}-src/configure" \
+      --prefix=/"${TARGET}/usr" \
+      --with-sysroot="/${TARGET}" \
+      --disable-bootstrap \
+      --disable-libmpx \
+      --disable-libsanitizer \
+      --disable-libunwind-exceptions \
+      --disable-libunwind-exceptions \
+      --disable-lto \
+      --disable-multilib \
+      --enable-__cxa_atexit \
+      --enable-gnu-indirect-function \
+      --enable-gnu-unique-object \
+      --enable-initfini-array \
+      --enable-languages="c,c++" \
+      --enable-linker-build-id \
+      --enable-plugin \
+      --enable-shared \
+      --enable-threads=posix \
+      --with-default-libstdcxx-abi=${LIBSTDCXX_ABI} \
+      --with-gcc-major-version-only \
+      --with-linker-hash-style="gnu" \
+      --with-tune="generic" \
+      && \
+    make -j 42 && \
+    make install
+
+
+# Create the devtoolset libstdc++ linkerscript that links dynamically against
+# the system libstdc++ 4.4 and provides all other symbols statically.
+case "${VERSION}" in
+devtoolset-7)
+mv "/${TARGET}/usr/lib/libstdc++.so.${LIBSTDCXX_VERSION}" \
+   "/${TARGET}/usr/lib/libstdc++.so.${LIBSTDCXX_VERSION}.backup"
+echo -e "OUTPUT_FORMAT(elf64-x86-64)\nINPUT ( libstdc++.so.6.0.13 -lstdc++_nonshared44 )" \
+   > "/${TARGET}/usr/lib/libstdc++.so.${LIBSTDCXX_VERSION}"
+cp "./x86_64-pc-linux-gnu/libstdc++-v3/src/.libs/libstdc++_nonshared44.a" \
+   "/${TARGET}/usr/lib"
+  ;;
+devtoolset-9)
+# Note that the installation path for libstdc++ here is /${TARGET}/usr/lib64/
+mv "/${TARGET}/usr/lib64/libstdc++.so.${LIBSTDCXX_VERSION}" \
+   "/${TARGET}/usr/lib64/libstdc++.so.${LIBSTDCXX_VERSION}.backup"
+echo -e "OUTPUT_FORMAT(elf64-x86-64)\nINPUT ( libstdc++.so.6.0.18 -lstdc++_nonshared44 )" \
+   > "/${TARGET}/usr/lib64/libstdc++.so.${LIBSTDCXX_VERSION}"
+cp "./x86_64-pc-linux-gnu/libstdc++-v3/src/.libs/libstdc++_nonshared44.a" \
+   "/${TARGET}/usr/lib64"
+;;
+esac
+
+# Link in architecture specific includes from the system; note that we cannot
+# link in the whole x86_64-linux-gnu folder, as otherwise we're overlaying
+# system gcc paths that we do not want to find.
+# TODO(klimek): Automate linking in all non-gcc / non-kernel include
+# directories.
+mkdir -p "/${TARGET}/usr/include/x86_64-linux-gnu"
+PYTHON_VERSIONS=("python3.10" "python3.11" "python3.12")
+for v in "${PYTHON_VERSIONS[@]}"; do
+  ln -s "/usr/local/include/${v}" "/${TARGET}/usr/include/x86_64-linux-gnu/${v}"
+done
+
+# Patch glibc to be compatable with modern clang
+case "${VERSION}" in
+devtoolset-9)
+  cd /
+  patch -p0 < /glibc2.17-inline.patch
+;;
+esac

ci/official/containers/ml_build/builder.devtoolset/fixlinks.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Re-direct all links in $1 that point to /lib... to point to $1/lib... instead.
+
+BASE="$1"
+find "${BASE}" -type l | \
+  while read l ; do
+    if [[ "$(readlink "$l")" == /lib* ]]; then
+      ORIG="$(readlink "$l")";
+      rm "$l";
+      ln -s "${BASE}${ORIG}" "$l"
+    fi
+  done
+

ci/official/containers/ml_build/builder.devtoolset/glibc2.17-inline.patch

@@ -0,0 +1,11 @@
+--- /dt9/usr/include/x86_64-linux-gnu/sys/cdefs.h       2013-09-30 13:58:17.000000000 +0000
++++ /dt9/usr/include/x86_64-linux-gnu/sys/cdefs.new.h   2022-11-04 17:17:31.727061220 +0000
+@@ -320,7 +320,7 @@
+
+ /* GCC 4.3 and above with -std=c99 or -std=gnu99 implements ISO C99
+    inline semantics, unless -fgnu89-inline is used.  */
+-#if (!defined __cplusplus || __GNUC_PREREQ (4,3)) && defined __GNUC__
++#if (!defined __cplusplus || __GNUC_PREREQ (4,3) || defined __clang__) && defined __GNUC__
+ # if defined __GNUC_STDC_INLINE__ || defined __cplusplus
+ #  define __extern_inline extern __inline __attribute__ ((__gnu_inline__))
+ #  define __extern_always_inline \

ci/official/containers/ml_build/builder.devtoolset/rpm-patch.sh

@@ -0,0 +1,28 @@
+#!/bin/bash -eu
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Given an RPM spec file $1, apply its patches.
+
+SPEC="$1"
+grep '%patch' "${SPEC}" |while read cmd ; do
+  N=$(echo "${cmd}" |sed 's,%patch\([0-9]\+\).*,\1,')
+  file=$(grep "Patch$N:" "${SPEC}" |sed 's,.*: ,,')
+  parg=$(echo "${cmd}" |sed 's,.*\(-p[0-9]\).*,\1,')
+  if [[ ! "${file}" =~ doxygen && "${cmd}" != \#* ]]; then
+    echo "patch ${parg} -s < ${file}"
+    patch ${parg} -s < "${file}"
+  fi
+done

ci/official/containers/ml_build/builder.packages.txt

@@ -0,0 +1,18 @@
+# Other build-related tools
+apt-transport-https
+autoconf
+automake
+build-essential
+ca-certificates
+curl
+git
+parallel
+sudo
+swig
+unzip
+zip
+openjdk-25-jdk
+vim
+wget
+jq
+file

ci/official/containers/ml_build/builder.requirements.txt

@@ -0,0 +1,20 @@
+# For wheel verification, and uploading
+auditwheel ~= 6.1.0
+twine ~= 6.1.0
+id
+urllib3
+requests
+
+# For XLA
+pyyaml
+
+# For JAX
+build ~= 1.2.2
+# uv is faster than pip for installing Python packages.
+uv ~= 0.5.30
+
+# For running wheel verification script
+immutabledict ~= 4.2.1
+
+# For MLIR's stubgen.py script
+typing-extensions

ci/official/containers/ml_build/cmake_toolchain.packages.txt

@@ -0,0 +1,8 @@
+cmake
+ninja-build
+ccache
+clang-18
+lld-18
+libclang-rt-18-dev
+libstdc++-13-dev
+lsb-release

ci/official/containers/ml_build/cuda12.1_cudnn9.1.packages.txt

@@ -0,0 +1,21 @@
+# All required CUDA packages
+cuda-compat-12-1
+cuda-command-line-tools-12-1
+cuda-cudart-dev-12-1
+cuda-nvcc-12-1
+cuda-cupti-12-1
+cuda-nvprune-12-1
+cuda-libraries-12-1
+cuda-libraries-dev-12-1
+cuda-nvml-dev-12-1
+libcufft-12-1
+libcurand-12-1
+libcusolver-dev-12-1
+libcusparse-dev-12-1
+libcublas-12-1
+libcublas-dev-12-1
+libnccl-dev=2.18.3-1+cuda12.1
+libnccl2=2.18.3-1+cuda12.1
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-dev-cuda-12=9.1.1.17-1
+libcudnn9-cuda-12=9.1.1.17-1

ci/official/containers/ml_build/cuda12.1_cudnn9.8.packages.txt

@@ -0,0 +1,21 @@
+# All required CUDA packages
+cuda-compat-12-1
+cuda-command-line-tools-12-1
+cuda-cudart-dev-12-1
+cuda-nvcc-12-1
+cuda-cupti-12-1
+cuda-nvprune-12-1
+cuda-libraries-12-1
+cuda-libraries-dev-12-1
+cuda-nvml-dev-12-1
+libcufft-12-1
+libcurand-12-1
+libcusolver-dev-12-1
+libcusparse-dev-12-1
+libcublas-12-1
+libcublas-dev-12-1
+libnccl-dev=2.18.3-1+cuda12.1
+libnccl2=2.18.3-1+cuda12.1
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-dev-cuda-12=9.8.0.87-1
+libcudnn9-cuda-12=9.8.0.87-1

ci/official/containers/ml_build/cuda12.3_cudnn9.1.packages.txt

@@ -0,0 +1,21 @@
+# All required CUDA packages
+cuda-compat-12-3
+cuda-command-line-tools-12-3
+cuda-cudart-dev-12-3
+cuda-nvcc-12-3
+cuda-cupti-12-3
+cuda-nvprune-12-3
+cuda-libraries-12-3
+cuda-libraries-dev-12-3
+cuda-nvml-dev-12-3
+libcufft-12-3
+libcurand-12-3
+libcusolver-dev-12-3
+libcusparse-dev-12-3
+libcublas-12-3
+libcublas-dev-12-3
+libnccl-dev=2.19.3-1+cuda12.3
+libnccl2=2.19.3-1+cuda12.3
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-dev-cuda-12=9.1.1.17-1
+libcudnn9-cuda-12=9.1.1.17-1

ci/official/containers/ml_build/cuda12.8_cudnn9.8.packages.txt

@@ -0,0 +1,21 @@
+# All required CUDA packages
+cuda-compat-12-8
+cuda-command-line-tools-12-8
+cuda-cudart-dev-12-8
+cuda-nvcc-12-8
+cuda-cupti-12-8
+cuda-nvprune-12-8
+cuda-libraries-12-8
+cuda-libraries-dev-12-8
+cuda-nvml-dev-12-8
+libcufft-12-8
+libcurand-12-8
+libcusolver-dev-12-8
+libcusparse-dev-12-8
+libcublas-12-8
+libcublas-dev-12-8
+libnccl-dev=2.25.1-1+cuda12.8
+libnccl2=2.25.1-1+cuda12.8
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-dev-cuda-12=9.8.0.87-1
+libcudnn9-cuda-12=9.8.0.87-1

ci/official/containers/ml_build/cuda13.0_cudnn9.12.packages.txt

@@ -0,0 +1,23 @@
+# All required CUDA packages
+cuda-compat-13-0
+cuda-command-line-tools-13-0
+cuda-cudart-dev-13-0
+cuda-nvcc-13-0
+cuda-cupti-13-0
+cuda-nvprune-13-0
+cuda-libraries-13-0
+cuda-libraries-dev-13-0
+cuda-nvml-dev-13-0
+libcufft-13-0
+libcurand-13-0
+libcusolver-dev-13-0
+libcusparse-dev-13-0
+libcublas-13-0
+libcublas-dev-13-0
+libnccl-dev=2.27.7-1+cuda13.0
+libnccl2=2.27.7-1+cuda13.0
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-headers-cuda-13=9.12.0.46-1
+libcudnn9-static-cuda-13=9.12.0.46-1
+libcudnn9-dev-cuda-13=9.12.0.46-1
+libcudnn9-cuda-13=9.12.0.46-1

ci/official/containers/ml_build/cuda13.0_cudnn9.15.packages.txt

@@ -0,0 +1,23 @@
+# All required CUDA packages
+cuda-compat-13-0
+cuda-command-line-tools-13-0
+cuda-cudart-dev-13-0
+cuda-nvcc-13-0
+cuda-cupti-13-0
+cuda-nvprune-13-0
+cuda-libraries-13-0
+cuda-libraries-dev-13-0
+cuda-nvml-dev-13-0
+libcufft-13-0
+libcurand-13-0
+libcusolver-dev-13-0
+libcusparse-dev-13-0
+libcublas-13-0
+libcublas-dev-13-0
+libnccl-dev=2.27.7-1+cuda13.0
+libnccl2=2.27.7-1+cuda13.0
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-headers-cuda-13=9.15.1.9-1
+libcudnn9-static-cuda-13=9.15.1.9-1
+libcudnn9-dev-cuda-13=9.15.1.9-1
+libcudnn9-cuda-13=9.15.1.9-1

ci/official/containers/ml_build/cuda13.2_cudnn9.15.packages.txt

@@ -0,0 +1,23 @@
+# All required CUDA packages
+cuda-compat-13-2
+cuda-command-line-tools-13-2
+cuda-cudart-dev-13-2
+cuda-nvcc-13-2
+cuda-cupti-13-2
+cuda-nvprune-13-2
+cuda-libraries-13-2
+cuda-libraries-dev-13-2
+cuda-nvml-dev-13-2
+libcufft-13-2
+libcurand-13-2
+libcusolver-dev-13-2
+libcusparse-dev-13-2
+libcublas-13-2
+libcublas-dev-13-2
+libnccl-dev=2.29.7-1+cuda13.2
+libnccl2=2.29.7-1+cuda13.2
+# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
+libcudnn9-headers-cuda-13=9.15.1.9-1
+libcudnn9-static-cuda-13=9.15.1.9-1
+libcudnn9-dev-cuda-13=9.15.1.9-1
+libcudnn9-cuda-13=9.15.1.9-1

ci/official/containers/ml_build/rbe_nvidia.packages.txt

@@ -0,0 +1,7 @@
+# The RBE machine itself has older kernel mode driver, and it requires
+# nvidia driver to be installed.
+nvidia-driver-580-open
+# TODO(b/445248346): The Docker image shouldn't have cuda-compat installed.
+# However, hermetic CUDA forward-compatibility mode is still missing some
+# libraries.
+cuda-compat-13-0

ci/official/containers/ml_build/setup.packages.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+#
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# setup.packages.sh: Given a list of Ubuntu packages, install them and clean up.
+# Usage: setup.packages.sh <package_list.txt>
+set -e
+
+# Prevent apt install tzinfo from asking our location (assumes UTC)
+export DEBIAN_FRONTEND=noninteractive
+
+# Update index for latest packages and upgrade base image packages.
+# Updating base image packages pulls in security fixes.
+apt-get update && apt-get upgrade -y
+# Remove commented lines and blank lines
+apt-get install -y --no-install-recommends $(sed -e '/^\s*#.*$/d' -e '/^\s*$/d' "$1" | sort -u)
+rm -rf /var/lib/apt/lists/*

ci/official/containers/ml_build/setup.python.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+#
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# setup.python.sh: Install a specific Python version and packages for it.
+# Usage: setup.python.sh <pyversion> <requirements.txt>
+set -xe
+
+source ~/.bashrc
+VERSION=$1
+REQUIREMENTS=$2
+
+# Install Python packages for this container's version
+if [[ ${VERSION} == "python3.13-nogil" || ${VERSION} == "python3.14-nogil" ]]; then
+  cat >pythons.txt <<EOF
+$VERSION
+EOF
+elif [[ ${VERSION} == "python3.14" || ${VERSION} == "python3.13" || ${VERSION} == "python3.12" ]]; then
+  cat >pythons.txt <<EOF
+$VERSION
+$VERSION-dev
+$VERSION-venv
+EOF
+else
+  cat >pythons.txt <<EOF
+$VERSION
+$VERSION-dev
+$VERSION-venv
+$VERSION-distutils
+EOF
+fi
+
+/setup.packages.sh pythons.txt
+
+# Python 3.10 include headers fix:
+# sysconfig.get_path('include') incorrectly points to /usr/local/include/python
+# map /usr/include/python3.10 to /usr/local/include/python3.10
+if [[ ! -f "/usr/local/include/$VERSION" ]]; then
+  ln -sf /usr/include/$VERSION /usr/local/include/$VERSION
+fi
+
+# Install pip
+wget --retry-connrefused --waitretry=1 --read-timeout=20 --timeout=15 --tries=5 https://bootstrap.pypa.io/get-pip.py
+/usr/bin/$VERSION get-pip.py
+/usr/bin/$VERSION -m pip install --no-cache-dir --upgrade pip
+/usr/bin/$VERSION -m pip install -U setuptools
+
+
+# For Python 3.13t, do not install twine as it does not have pre-built wheels
+# for this Python version and building it from source fails. We only need twine
+# to be present on the system Python which in this case is 3.12.
+# Same reason for Python 3.14.
+if [[ ${VERSION} == "python3.13-nogil" || ${VERSION} == "python3.14" || ${VERSION} == "python3.14-nogil" ]]; then
+  grep -v "twine" $REQUIREMENTS > requirements_without_twine.txt
+  REQUIREMENTS=requirements_without_twine.txt
+fi
+
+# Disable the cache dir to save image space, and install packages
+/usr/bin/$VERSION -m pip install --no-cache-dir -r $REQUIREMENTS -U

ci/official/containers/ml_build/setup.sources.cudnn.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+#
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Sets up custom apt sources for our TF images.
+
+# Prevent apt install tzinfo from asking our location (assumes UTC)
+export DEBIAN_FRONTEND=noninteractive
+
+# Fetch the NVIDIA key.
+apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub;
+
+# Set up sources for NVIDIA CUDNN.
+cat >/etc/apt/sources.list.d/nvidia.list <<SOURCES
+# NVIDIA
+deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /
+
+SOURCES

ci/official/containers/ml_build/setup.sources.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+#
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Sets up custom apt sources for our TF images.
+
+# Prevent apt install tzinfo from asking our location (assumes UTC)
+export DEBIAN_FRONTEND=noninteractive
+
+# Set up shared custom sources
+apt-get update
+apt-get install -y gnupg ca-certificates software-properties-common
+add-apt-repository -y universe
+apt-get update
+
+# Deadsnakes: https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa
+apt-key adv --keyserver keyserver.ubuntu.com --recv-keys F23C5A6CF475977595C89F51BA6932366A755776
+
+# LLVM/Clang: https://apt.llvm.org/
+apt-key adv --fetch-keys https://apt.llvm.org/llvm-snapshot.gpg.key
+
+# Set up custom sources
+cat >/etc/apt/sources.list.d/custom.list <<SOURCES
+# More Python versions: Deadsnakes
+deb http://ppa.launchpad.net/deadsnakes/ppa/ubuntu jammy main
+deb-src http://ppa.launchpad.net/deadsnakes/ppa/ubuntu jammy main
+
+# LLVM/Clang repository
+deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main
+deb-src http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main
+SOURCES

ci/official/containers/ml_build/setup.toolchain.cmake.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+#
+# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Sets up custom apt sources for our TF images.
+
+# Prevent apt install tzinfo from asking our location (assumes UTC)
+export DEBIAN_FRONTEND=noninteractive
+
+# Install `software-properties-common` for `add-apt-repository`.
+apt-get update
+apt-get install -y software-properties-common
+
+# Add Ubuntu Toolchain PPA for newer GCC/libstdc++.
+add-apt-repository -y ppa:ubuntu-toolchain-r/test
+
+# Set up symlinks for Clang and LLD.
+ln -sf /usr/bin/clang-18 /usr/bin/clang
+ln -sf /usr/bin/clang++-18 /usr/bin/clang++
+ln -sf /usr/bin/lld-18 /usr/bin/lld