Merge pull request #2287 from ROCmSoftwarePlatform/develop-upstream-s…

…ync-231030

Develop upstream sync 231030

.bazelrc

@@ -56,6 +56,7 @@
 #
 #     rbe_linux_cpu:                  RBE options to build with only CPU support.
 #     rbe_linux_cuda:                 RBE options to build with GPU support using clang.
+#     rbe_linux_cuda_nvcc:            RBE options to build with GPU support using nvcc.
 #
 #     rbe_win_py39: Windows Python 3.9 RBE config
 #
@@ -238,9 +239,12 @@ build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
 # Select supported compute capabilities (supported graphics cards).
 # This is the same as the official TensorFlow builds.
 # See https://developer.nvidia.com/cuda-gpus#compute
-# TODO(angerson, perfinion): What does sm_ vs compute_ mean? How can users
-# select a good value for this? See go/tf-pip-cuda
-build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_75,compute_80"
+# `compute_XY` enables PTX embedding in addition to SASS. PTX
+# is forward compatible beyond the current compute capability major
+# release while SASS is only forward compatible inside the current
+# major release. Example: sm_80 kernels can run on sm_89 GPUs but
+# not on sm_90 GPUs. compute_80 kernels though can also run on sm_90 GPUs.
+build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
 
 # Set up compilation CUDA version and paths and use the CUDA Clang toolchain.
 build:cuda_clang_official --config=cuda_clang
@@ -250,7 +254,7 @@ build:cuda_clang_official --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-12.2"
 build:cuda_clang_official --action_env=GCC_HOST_COMPILER_PATH="/dt9/usr/bin/gcc"
 build:cuda_clang_official --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-17/bin/clang"
 build:cuda_clang_official --action_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
-build:cuda_clang_official --crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain"
+build:cuda_clang_official --crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain"
 
 # Debug config
 build:dbg -c dbg
@@ -488,12 +492,12 @@ build:rbe_linux --host_linkopt=-lm
 
 build:rbe_linux_cpu --config=rbe_linux
 # Linux cpu and cuda builds share the same toolchain now.
-build:rbe_linux_cpu --host_crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain"
-build:rbe_linux_cpu --crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain"
-build:rbe_linux_cpu --extra_toolchains="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain-linux-x86_64"
-build:rbe_linux_cpu --extra_execution_platforms="@sigbuild-r2.14-clang_config_platform//:platform"
-build:rbe_linux_cpu --host_platform="@sigbuild-r2.14-clang_config_platform//:platform"
-build:rbe_linux_cpu --platforms="@sigbuild-r2.14-clang_config_platform//:platform"
+build:rbe_linux_cpu --host_crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain"
+build:rbe_linux_cpu --crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain"
+build:rbe_linux_cpu --extra_toolchains="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain-linux-x86_64"
+build:rbe_linux_cpu --extra_execution_platforms="@sigbuild-r2.16-clang_config_platform//:platform"
+build:rbe_linux_cpu --host_platform="@sigbuild-r2.16-clang_config_platform//:platform"
+build:rbe_linux_cpu --platforms="@sigbuild-r2.16-clang_config_platform//:platform"
 # This is needed for all Clang17 builds but must not be present in GCC builds.
 build:rbe_linux_cpu --copt=-Wno-error=unused-command-line-argument
 # This was added in clang-16 by https://reviews.llvm.org/D133574.
@@ -502,7 +506,7 @@ build:rbe_linux_cpu --copt=-Wno-error=unused-command-line-argument
 # See https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183.
 build:rbe_linux_cpu --copt=-Wno-gnu-offsetof-extensions
 # Python config is the same across all containers because the binary is the same
-build:rbe_linux_cpu --repo_env=TF_PYTHON_CONFIG_REPO="@sigbuild-r2.14-clang_config_python"
+build:rbe_linux_cpu --repo_env=TF_PYTHON_CONFIG_REPO="@sigbuild-r2.16-clang_config_python"
 build:rbe_linux_cpu --python_path="/usr/bin/python3"
 # These you may need to change for your own GCP project.
 common:rbe_linux_cpu --remote_instance_name=projects/tensorflow-testing/instances/default_instance
@@ -523,9 +527,9 @@ build:rbe_linux_cuda --config=cuda_clang_official
 build:rbe_linux_cuda --config=rbe_linux_cpu
 # For Remote build execution -- GPU configuration
 build:rbe_linux_cuda --repo_env=REMOTE_GPU_TESTING=1
-build:rbe_linux_cuda --repo_env=TF_CUDA_CONFIG_REPO="@sigbuild-r2.14-clang_config_cuda"
-build:rbe_linux_cuda --repo_env=TF_TENSORRT_CONFIG_REPO="@sigbuild-r2.14-clang_config_tensorrt"
-build:rbe_linux_cuda --repo_env=TF_NCCL_CONFIG_REPO="@sigbuild-r2.14-clang_config_nccl"
+build:rbe_linux_cuda --repo_env=TF_CUDA_CONFIG_REPO="@sigbuild-r2.16-clang_config_cuda"
+build:rbe_linux_cuda --repo_env=TF_TENSORRT_CONFIG_REPO="@sigbuild-r2.16-clang_config_tensorrt"
+build:rbe_linux_cuda --repo_env=TF_NCCL_CONFIG_REPO="@sigbuild-r2.16-clang_config_nccl"
 test:rbe_linux_cuda --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
 
 # ROCm
@@ -539,6 +543,35 @@ build:rbe_linux_rocm_base --platforms="@ubuntu20.04-gcc9_manylinux2014-rocm_conf
 build:rbe_linux_rocm_base --action_env=TF_ROCM_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-rocm_config_rocm"
 build:rbe_linux_rocm_py3.9 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-rocm_config_python3.9"
 
+build:rbe_linux_cuda_nvcc --config=cuda
+build:rbe_linux_cuda_nvcc --repo_env TF_NCCL_USE_STUB=1
+build:rbe_linux_cuda_nvcc --@local_xla//xla/python:enable_gpu=true
+build:rbe_linux_cuda_nvcc --@local_xla//xla/python:jax_cuda_pip_rpaths=true
+build:rbe_linux_cuda_nvcc --define=xla_python_enable_gpu=true
+build:rbe_linux_cuda_nvcc --config=tensorrt
+build:rbe_linux_cuda_nvcc --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_75,compute_80"
+build:rbe_linux_cuda_nvcc --action_env=TF_CUDA_VERSION="12"
+build:rbe_linux_cuda_nvcc --action_env=TF_CUDNN_VERSION="8"
+build:rbe_linux_cuda_nvcc --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-12.2"
+build:rbe_linux_cuda_nvcc --action_env=GCC_HOST_COMPILER_PATH="/dt9/usr/bin/gcc"
+build:rbe_linux_cuda_nvcc --action_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
+build:rbe_linux_cuda_nvcc --crosstool_top="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain"
+build:rbe_linux_cuda_nvcc --config=rbe_linux
+build:rbe_linux_cuda_nvcc --host_crosstool_top="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain"
+build:rbe_linux_cuda_nvcc --extra_toolchains="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain-linux-x86_64"
+build:rbe_linux_cuda_nvcc --extra_execution_platforms="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform"
+build:rbe_linux_cuda_nvcc --host_platform="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform"
+build:rbe_linux_cuda_nvcc --platforms="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform"
+build:rbe_linux_cuda_nvcc --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_python3.9"
+build:rbe_linux_cuda_nvcc --python_path="/usr/bin/python3"
+# These you may need to change for your own GCP project.
+common:rbe_linux_cuda_nvcc --remote_instance_name=projects/tensorflow-testing/instances/default_instance
+build:rbe_linux_cuda_nvcc --repo_env=REMOTE_GPU_TESTING=1
+build:rbe_linux_cuda_nvcc --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_cuda"
+build:rbe_linux_cuda_nvcc --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_tensorrt"
+build:rbe_linux_cuda_nvcc --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_nccl"
+test:rbe_linux_cuda_nvcc --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
+
 # TODO(kanglan): Remove rbe_win and rbe_win_py3* after b/289091160 is fixed
 build:rbe_win --config=rbe_base
 build:rbe_win --crosstool_top="//tensorflow/tools/toolchains/win/tf_win_05022023:toolchain"
@@ -593,8 +626,6 @@ try-import %workspace%/.bazelrc.user
 # Here are bazelrc configs for release builds
 # Build TensorFlow v2.
 test:release_base --test_size_filters=small,medium
-# TODO(b/294367488) disable after 2.15 brancut
-test:release_base --flaky_test_attempts=3
 
 # Target the AVX instruction set
 build:release_linux_base --config=avx_linux
@@ -632,7 +663,7 @@ test:release_linux_base --test_summary=short
 
 # Use the Clang toolchain to compile
 build:release_cpu_linux --config=release_linux_base
-build:release_cpu_linux --crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain"
+build:release_cpu_linux --crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain"
 
 build:release_gpu_linux --config=release_cpu_linux
 # Set up compilation CUDA version and paths and use the CUDA Clang toolchain.

.github/workflows/osv-scanner-scheduled.yml

@@ -0,0 +1,39 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+name: OSV-Scanner Scheduled Scan
+
+on:
+  schedule:
+    - cron: 0 4 * * 1
+
+permissions:
+  # Require writing security events to upload SARIF file to security tab
+  security-events: write
+  # Only need to read contents
+  contents: read
+
+jobs:
+  scan-scheduled:
+    uses: "google/osv-scanner/.github/workflows/osv-scanner-reusable.yml@main"
+    with:
+      scan-args: |-
+        --lockfile=requirements.txt:./requirements_lock_3_9.txt
+        --lockfile=requirements.txt:./requirements_lock_3_10.txt
+        --lockfile=requirements.txt:./requirements_lock_3_11.txt
+        --lockfile=requirements.txt:./requirements_lock_3_12.txt
+        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.requirements.txt
+        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/jax.requirements.txt
+        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.usertools/test.requirements.txt

.github/workflows/sigbuild-docker-branch.yml

@@ -34,7 +34,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [python3.9, python3.10, python3.11]
+        python-version: [python3.9, python3.10, python3.11, python3.12]
     steps:
       - name: Delete unnecessary tools folder
         run: rm -rf /opt/hostedtoolcache

.github/workflows/sigbuild-docker-presubmit.yml

@@ -32,7 +32,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [python3.9, python3.10, python3.11]
+        python-version: [python3.9, python3.10, python3.11, python3.12]
     permissions:
       contents: read
       pull-requests: write
@@ -87,6 +87,7 @@ jobs:
           message: |
             I pushed these containers:
             
+            - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.12`
             - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.11`
             - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.10`
             - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.9`

.github/workflows/sigbuild-docker.yml

@@ -37,7 +37,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [python3.9, python3.10, python3.11]
+        python-version: [python3.9, python3.10, python3.11, python3.12]
     steps:
       - name: Delete unnecessary tools folder
         run: rm -rf /opt/hostedtoolcache

.github/workflows/update-rbe.yml

@@ -105,6 +105,18 @@ jobs:
         map sigbuild-r2.14-clang-python3.9 2.14-python3.9
         map sigbuild-r2.14-clang-python3.10 2.14-python3.10
         map sigbuild-r2.14-clang-python3.11 2.14-python3.11
+        # TF 2.16
+        map sigbuild-r2.16 2.16-python3.9
+        map sigbuild-r2.16-python3.9 2.16-python3.9
+        map sigbuild-r2.16-python3.10 2.16-python3.10
+        map sigbuild-r2.16-python3.11 2.16-python3.11
+        map sigbuild-r2.16-python3.12 2.16-python3.12
+        # TF 2.16 + Clang (containers are the same, but env vars in configs.bzl are different)
+        map sigbuild-r2.16-clang 2.16-python3.9
+        map sigbuild-r2.16-clang-python3.9 2.16-python3.9
+        map sigbuild-r2.16-clang-python3.10 2.16-python3.10
+        map sigbuild-r2.16-clang-python3.11 2.16-python3.11
+        map sigbuild-r2.16-clang-python3.12 2.16-python3.12
     - name: Create Pull Request with changes
       uses: peter-evans/create-pull-request@2b011faafdcbc9ceb11414d64d0573f37c774b04 # v4.2.3
       with:

RELEASE.md

@@ -28,6 +28,8 @@
 
 * `tf.lite`
     * Added support for `stablehlo.gather`.
+    * Added support for `stablehlo.add`.
+    * Added support for `stablehlo.multiply`.
 
 ## Keras
 

ci/official/containers/linux_arm64/devel.packages.txt

@@ -3,6 +3,8 @@ autoconf
 automake
 build-essential
 ca-certificates
+# TODO(b/308399490) Remove CMake once dm-tree (Keras dependency) has 3.12 wheels
+cmake
 llvm-17
 clang-17
 clang-format-12

ci/official/envs/ci_default

@@ -16,6 +16,7 @@ TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
 TFCI_NVIDIA_SMI_ENABLE=
 TFCI_OUTPUT_DIR=build_output
 TFCI_LIBTPU_DOWNLOAD_ENABLE=0
+TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=0
 TFCI_LIBTPU_DOWNLOAD_URL=
 TFCI_UPLOAD_LIB_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_ENABLE=

ci/official/envs/ci_nightly_uploads

@@ -1,6 +1,8 @@
-TFCI_UPLOAD_LIB_ENABLE=0
-TFCI_UPLOAD_LIB_LATEST_ENABLE=
-TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_LIB_ENABLE=1
+TFCI_UPLOAD_LIB_GCS_URI="gs://libtensorflow-nightly/$(date -I)"
+TFCI_UPLOAD_LIB_LATEST_ENABLE=1
+TFCI_UPLOAD_LIB_LATEST_GCS_URI="gs://libtensorflow-nightly/latest"
+TFCI_UPLOAD_WHL_GCS_ENABLE=0
 TFCI_UPLOAD_WHL_GCS_URI=
-TFCI_UPLOAD_WHL_PYPI_ARGS=
-TFCI_UPLOAD_WHL_PYPI_ENABLE=
+TFCI_UPLOAD_WHL_PYPI_ARGS=(--config-file="$KOKORO_KEYSTORE_DIR/73361_tensorflow_pypirc_using_global_api_token" --repository pypi-warehouse)
+TFCI_UPLOAD_WHL_PYPI_ENABLE=1

ci/official/envs/disable_all_uploads

@@ -1,7 +1,9 @@
-TFCI_UPLOAD_LIB_ENABLE=0
+TFCI_DOCKER_REBUILD_UPLOAD_ENABLE=0
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_GCS_URI=
 TFCI_UPLOAD_LIB_LATEST_ENABLE=
-TFCI_DOCKER_REBUILD_UPLOAD_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_GCS_URI=
 TFCI_UPLOAD_WHL_GCS_ENABLE=
 TFCI_UPLOAD_WHL_GCS_URI=
-TFCI_UPLOAD_WHL_PYPI_ENABLE=
 TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/nightly_linux_arm64_cpu_py310

@@ -1,5 +1,6 @@
 source ci/official/envs/ci_default
-source ci/official/envs/ci_nightly_uploads
+# Disable arm64 uploads while being worked on
+source ci/official/envs/disable_all_uploads
 TFCI_PYTHON_VERSION=3.10
 TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64

ci/official/envs/nightly_linux_arm64_cpu_py311

@@ -1,5 +1,6 @@
 source ci/official/envs/ci_default
-source ci/official/envs/ci_nightly_uploads
+# Disable arm64 uploads while being worked on
+source ci/official/envs/disable_all_uploads
 TFCI_PYTHON_VERSION=3.11
 TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64

ci/official/envs/nightly_linux_arm64_cpu_py312

@@ -0,0 +1,10 @@
+source ci/official/envs/ci_default
+# Disable arm64 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_PYTHON_VERSION=3.12
+TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag)
+TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python
+TFCI_DOCKER_REBUILD_ARGS=(--target=tf ci/official/containers/linux_arm64)
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1

ci/official/envs/nightly_linux_arm64_cpu_py39

@@ -1,5 +1,6 @@
 source ci/official/envs/ci_default
-source ci/official/envs/ci_nightly_uploads
+# Disable arm64 uploads while being worked on
+source ci/official/envs/disable_all_uploads
 TFCI_PYTHON_VERSION=3.9
 TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64

ci/official/envs/nightly_linux_x86_cpu_py312

@@ -0,0 +1,10 @@
+source ci/official/envs/ci_default
+# Disable 3.12 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_PYTHON_VERSION=3.12
+TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cpu
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION}
+TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles)
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1

ci/official/envs/nightly_linux_x86_cuda_py312

@@ -0,0 +1,11 @@
+source ci/official/envs/ci_default
+# Disable 3.12 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_PYTHON_VERSION=3.12
+TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--nightly_flag)
+TFCI_DOCKER_ARGS=(--gpus all)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION}
+TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles)
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1

ci/official/envs/nightly_linux_x86_tpu_py310

@@ -1,11 +1,11 @@
 source ci/official/envs/ci_default
-source ci/official/envs/ci_nightly_uploads
+# Disable tpu uploads while being worked on
+source ci/official/envs/disable_all_uploads
 TFCI_PYTHON_VERSION=3.10
 TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config=tpu)
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu
 TFCI_BUILD_PIP_PACKAGE_ARGS=(--tpu --nightly_flag)
 TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION}
 TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles)
-TFCI_LIBTPU_DOWNLOAD_ENABLE=1
-TFCI_LIBTPU_DOWNLOAD_URL=https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.8.0/libtpu.so
+TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=1
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1

ci/official/envs/nightly_linux_x86_tpu_py311

@@ -1,11 +1,11 @@
 source ci/official/envs/ci_default
-source ci/official/envs/ci_nightly_uploads
+# Disable tpu uploads while being worked on
+source ci/official/envs/disable_all_uploads
 TFCI_PYTHON_VERSION=3.11
 TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config=tpu)
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu
 TFCI_BUILD_PIP_PACKAGE_ARGS=(--tpu --nightly_flag)
 TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION}
 TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles)
-TFCI_LIBTPU_DOWNLOAD_ENABLE=1
-TFCI_LIBTPU_DOWNLOAD_URL=https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.8.0/libtpu.so
+TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=1
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1