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

…ync-231023 Develop upstream sync 231023
ROCm · Oct 27, 2023 · a1c584b · a1c584b
2 parents 8ef05e1 + ded7c3a
commit a1c584b
Show file tree

Hide file tree

Showing 464 changed files with 8,391 additions and 4,016 deletions.
diff --git a/.github/bot_config.yml b/.github/bot_config.yml
@@ -18,7 +18,7 @@ assignees:
    - sushreebarsa
    - SuryanarayanaY
    - tilakrayal
-   - Varsha-anjanappa
+   - Venkat6871
 # A list of assignees for compiler folder
 compiler_assignees:
    - joker-eph

diff --git a/RELEASE.md b/RELEASE.md
@@ -26,9 +26,24 @@
 * <IF A CHANGE CLOSES A GITHUB ISSUE, IT SHOULD BE DOCUMENTED HERE>
 * <NOTES SHOULD BE GROUPED PER AREA>
 
+* `tf.lite`
+    * Added support for `stablehlo.gather`.
+
 ## Keras
 
-<INSERT SMALL BLURB ABOUT RELEASE FOCUS AREA AND POTENTIAL TOOLCHAIN CHANGES>
+*  `keras.layers.experimental.DynamicEmbedding`
+    * Added `DynamicEmbedding` Keras layer
+    * Added 'UpdateEmbeddingCallback`
+    * `DynamicEmbedding` layer allows for the continuous updating of the
+      vocabulary and embeddings during the training process. This layer
+      maintains a hash table to track the most up-to-date vocabulary based on
+      the inputs received by the layer and the eviction policy. When this layer
+      is used with an `UpdateEmbeddingCallback`, which is a time-based callback,
+      the vocabulary lookup tensor is updated at the time interval set in the
+      `UpdateEmbeddingCallback` based on the most up-to-date vocabulary hash
+      table maintained by the layer. If this layer is not used in conjunction
+      with `UpdateEmbeddingCallback` the behavior of the layer would be same as
+      `keras.layers.Embedding`.
 
 ### Breaking Changes
 

diff --git a/WORKSPACE b/WORKSPACE
@@ -31,12 +31,12 @@ load(
 
 python_repository(name = "python_version_repo")
 
-load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION")
+load("@python_version_repo//:py_version.bzl", "TF_PYTHON_VERSION")
 
 python_register_toolchains(
     name = "python",
     ignore_root_user_error = True,
-    python_version = HERMETIC_PYTHON_VERSION,
+    python_version = TF_PYTHON_VERSION,
 )
 
 load("@python//:defs.bzl", "interpreter")
@@ -62,7 +62,7 @@ pip_parse(
     name = "pypi",
     annotations = NUMPY_ANNOTATIONS,
     python_interpreter_target = interpreter,
-    requirements = "//:requirements_lock_" + HERMETIC_PYTHON_VERSION.replace(".", "_") + ".txt",
+    requirements = "//:requirements_lock_" + TF_PYTHON_VERSION.replace(".", "_") + ".txt",
 )
 
 load("@pypi//:requirements.bzl", "install_deps")

diff --git a/ci/official/envs/ci_default b/ci/official/envs/ci_default
@@ -1,6 +1,7 @@
 TFCI_BAZEL_BAZELRC_ARGS=()
 TFCI_BAZEL_CONFIG_PREFIX=
 TFCI_BAZEL_COMMON_ARGS=()
+TFCI_PYTHON_VERSION=
 TFCI_BUILD_PIP_PACKAGE_ARGS=()
 TFCI_DOCKER_ARGS=()
 TFCI_DOCKER_ENABLE=1

diff --git a/ci/official/envs/continuous_linux_arm64_cpu_py310 b/ci/official/envs/continuous_linux_arm64_cpu_py310
@@ -1,6 +1,7 @@
 # This envrionment is experimental and should not yet be used for production jobs
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.10
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64
-TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python
 TFCI_DOCKER_REBUILD_ARGS=(--target=tf ci/official/containers/linux_arm64)
diff --git a/ci/official/envs/continuous_linux_arm64_cpu_py311 b/ci/official/envs/continuous_linux_arm64_cpu_py311
@@ -1,6 +1,7 @@
 # This envrionment is experimental and should not yet be used for production jobs
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.11
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64
-TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=3.11)
+TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python
 TFCI_DOCKER_REBUILD_ARGS=(--target=tf ci/official/containers/linux_arm64)
diff --git a/ci/official/envs/continuous_linux_arm64_cpu_py39 b/ci/official/envs/continuous_linux_arm64_cpu_py39
@@ -1,6 +1,7 @@
 # This envrionment is experimental and should not yet be used for production jobs
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.9
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64
-TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=3.9)
+TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python
 TFCI_DOCKER_REBUILD_ARGS=(--target=tf ci/official/containers/linux_arm64)
diff --git a/ci/official/envs/continuous_linux_x86_cpu_py310 b/ci/official/envs/continuous_linux_x86_cpu_py310
@@ -1,5 +1,6 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.10
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cpu
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config rbe_linux_cpu --repo_env=TF_PYTHON_VERSION=3.10)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config rbe_linux_cpu --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-pythonlatest-python${TFCI_PYTHON_VERSION}
+TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles)
diff --git a/ci/official/envs/continuous_linux_x86_cpu_py311 b/ci/official/envs/continuous_linux_x86_cpu_py311
@@ -1,5 +1,6 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.11
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cpu
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config rbe_linux_cpu --repo_env=TF_PYTHON_VERSION=3.11)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.11
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.11 --target=devel tools/tf_sig_build_dockerfiles)
+TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config rbe_linux_cpu --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+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)
diff --git a/ci/official/envs/continuous_linux_x86_cpu_py39 b/ci/official/envs/continuous_linux_x86_cpu_py39
@@ -1,5 +1,6 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.9
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cpu
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config rbe_linux_cpu --repo_env=TF_PYTHON_VERSION=3.9)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.9
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.11 --target=devel tools/tf_sig_build_dockerfiles)
+TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config rbe_linux_cpu --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+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)
diff --git a/ci/official/envs/continuous_linux_x86_cuda_py310 b/ci/official/envs/continuous_linux_x86_cuda_py310
@@ -1,7 +1,8 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.10
 TFCI_NVIDIA_SMI_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config rbe_linux_cuda --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config rbe_linux_cuda --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ARGS=(--gpus all)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+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)
diff --git a/ci/official/envs/continuous_linux_x86_cuda_py311 b/ci/official/envs/continuous_linux_x86_cuda_py311
@@ -1,7 +1,8 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.11
 TFCI_NVIDIA_SMI_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config rbe_linux_cuda --repo_env=TF_PYTHON_VERSION=3.11)
+TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config rbe_linux_cuda --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ARGS=(--gpus all)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.11
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.11 --target=devel tools/tf_sig_build_dockerfiles)
+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)
diff --git a/ci/official/envs/continuous_linux_x86_cuda_py39 b/ci/official/envs/continuous_linux_x86_cuda_py39
@@ -1,7 +1,8 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.9
 TFCI_NVIDIA_SMI_ENABLE=1
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config rbe_linux_cuda --repo_env=TF_PYTHON_VERSION=3.9)
+TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config rbe_linux_cuda --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ARGS=(--gpus all)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.9
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.9 --target=devel tools/tf_sig_build_dockerfiles)
+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)
diff --git a/ci/official/envs/continuous_macos_arm64_py310 b/ci/official/envs/continuous_macos_arm64_py310
@@ -1,4 +1,5 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.10
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64
-TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ENABLE=0
diff --git a/ci/official/envs/continuous_macos_arm64_py311 b/ci/official/envs/continuous_macos_arm64_py311
@@ -1,4 +1,5 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.11
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64
-TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --repo_env=TF_PYTHON_VERSION=3.11)
+TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ENABLE=0
diff --git a/ci/official/envs/continuous_macos_arm64_py39 b/ci/official/envs/continuous_macos_arm64_py39
@@ -1,4 +1,5 @@
 source ci/official/envs/ci_default
+TFCI_PYTHON_VERSION=3.9
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64
-TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --repo_env=TF_PYTHON_VERSION=3.11)
+TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ENABLE=0
diff --git a/ci/official/envs/nightly_libtensorflow_linux_x86_cpu b/ci/official/envs/nightly_libtensorflow_linux_x86_cpu
@@ -1,7 +1,8 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+TFCI_PYTHON_VERSION=3.10
+TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
+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_LIB_SUFFIX="-cpu-linux-x86_64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
diff --git a/ci/official/envs/nightly_libtensorflow_linux_x86_cuda b/ci/official/envs/nightly_libtensorflow_linux_x86_cuda
@@ -1,9 +1,10 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_PYTHON_VERSION=3.10
+TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ARGS=(--gpus all)
-TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+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_LIB_SUFFIX="-gpu-linux-x86_64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
 TFCI_NVIDIA_SMI_ENABLE=1
diff --git a/ci/official/envs/nightly_libtensorflow_macos_arm64 b/ci/official/envs/nightly_libtensorflow_macos_arm64
@@ -1,6 +1,7 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_PYTHON_VERSION=3.10
+TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION)
 TFCI_DOCKER_ENABLE=0
 TFCI_LIB_SUFFIX="-cpu-macos-arm64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1

diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py310 b/ci/official/envs/nightly_linux_arm64_cpu_py310
@@ -1,6 +1,7 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10)
+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
 TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag)
 TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python

diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py311 b/ci/official/envs/nightly_linux_arm64_cpu_py311
@@ -1,6 +1,7 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.11)
+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
 TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag)
 TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python

diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py39 b/ci/official/envs/nightly_linux_arm64_cpu_py39
@@ -1,6 +1,7 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.9)
+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
 TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag)
 TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python

diff --git a/ci/official/envs/nightly_linux_x86_cpu_py310 b/ci/official/envs/nightly_linux_x86_cpu_py310
@@ -1,8 +1,9 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_PYTHON_VERSION=3.10
+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-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+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
diff --git a/ci/official/envs/nightly_linux_x86_cpu_py311 b/ci/official/envs/nightly_linux_x86_cpu_py311
@@ -1,8 +1,9 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.11)
+TFCI_PYTHON_VERSION=3.11
+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-python3.11
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.11 --target=devel tools/tf_sig_build_dockerfiles)
+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
diff --git a/ci/official/envs/nightly_linux_x86_cpu_py39 b/ci/official/envs/nightly_linux_x86_cpu_py39
@@ -1,8 +1,9 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.9)
+TFCI_PYTHON_VERSION=3.9
+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-python3.9
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.9 --target=devel tools/tf_sig_build_dockerfiles)
+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
diff --git a/ci/official/envs/nightly_linux_x86_cuda_py310 b/ci/official/envs/nightly_linux_x86_cuda_py310
@@ -1,9 +1,10 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10)
+TFCI_PYTHON_VERSION=3.10
+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-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+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
diff --git a/ci/official/envs/nightly_linux_x86_cuda_py311 b/ci/official/envs/nightly_linux_x86_cuda_py311
@@ -1,9 +1,10 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.11)
+TFCI_PYTHON_VERSION=3.11
+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-python3.11
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.11 --target=devel tools/tf_sig_build_dockerfiles)
+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
diff --git a/ci/official/envs/nightly_linux_x86_cuda_py39 b/ci/official/envs/nightly_linux_x86_cuda_py39
@@ -1,9 +1,10 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.9)
+TFCI_PYTHON_VERSION=3.9
+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-python3.9
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.9 --target=devel tools/tf_sig_build_dockerfiles)
+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
diff --git a/ci/official/envs/nightly_linux_x86_tpu_py310 b/ci/official/envs/nightly_linux_x86_tpu_py310
@@ -1,10 +1,11 @@
 source ci/official/envs/ci_default
 source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=3.10 --config=tpu)
+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-python3.10
-TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=3.10 --target=devel tools/tf_sig_build_dockerfiles)
+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_NIGHTLY_UPDATE_VERSION_ENABLE=1