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

…ync-230731

Develop upstream sync 230731

.bazelrc

@@ -202,8 +202,6 @@ build:monolithic --define framework_shared_object=false
 build:monolithic --define tsl_protobuf_header_only=false
 build:monolithic --experimental_link_static_libraries_once=false  # b/229868128
 
-build:linux --define=build_with_onednn_v2=true
-
 # Please note that MKL on MacOS is still not supported.
 # If you would like to use a local MKL instead of downloading, please set the
 # environment variable "TF_MKL_ROOT" every time before build.
@@ -417,6 +415,8 @@ build:rbe --google_default_credentials
 build:rbe --bes_backend=buildeventservice.googleapis.com
 build:rbe --bes_results_url="https://source.cloud.google.com/results/invocations"
 build:rbe --bes_timeout=600s
+# TODO(b/290857564): Remove from mainline after 2.14 branch cut 
+build:rbe --flaky_test_attempts=3
 build:rbe --define=EXECUTOR=remote
 build:rbe --jobs=800
 build:rbe --remote_executor=grpcs://remotebuildexecution.googleapis.com
@@ -557,6 +557,8 @@ try-import %workspace%/.bazelrc.user
 # Here are bazelrc configs for release builds
 build:release_base --config=v2
 test:release_base --test_size_filters=small,medium
+# TODO(b/290857564): Remove from mainline after 2.14 branch cut 
+test:release_base --flaky_test_attempts=3
 
 build:release_cpu_linux --config=release_base
 build:release_cpu_linux --config=avx_linux

RELEASE.md

@@ -88,6 +88,9 @@
       relayouting a tensor according to the layout of another tensor.
     * Added `dtensor.get_default_mesh`, for retrieving the current default
       mesh under the dtensor context.
+    * \*fft\* ops now support dtensors with any layout. Fixed bug in 'fft2d/
+      fft3d', 'ifft2d/ifft3d', 'rfft2d/rfft3d', and 'irfft2d/irfft3d' for
+      sharded input.
 
 *  `tf.experimental.strict_mode`
     * Added a new API, `strict_mode`, which converts all deprecation warnings

ci/official/envs/continuous_linux_x86_cpu_py310

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cpu.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config rbe)
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu)
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=()
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
+TFCI_NVIDIA_SMI_ENABLE=
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI=
+TFCI_UPLOAD_LIB_URI=
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/continuous_linux_x86_cpu_py311

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cpu.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config rbe)
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu)
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=()
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.11
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
+TFCI_NVIDIA_SMI_ENABLE=
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI=
+TFCI_UPLOAD_LIB_URI=
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/continuous_linux_x86_cpu_py39

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cpu.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config rbe)
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu)
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=()
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.9
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
+TFCI_NVIDIA_SMI_ENABLE=
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI=
+TFCI_UPLOAD_LIB_URI=
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/continuous_linux_x86_cuda_py310

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cuda.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config rbe)
+TFCI_BUILD_PIP_PACKAGE_ARGS=()
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=(--gpus all)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-gpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
+TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI=
+TFCI_UPLOAD_LIB_URI=
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/continuous_linux_x86_cuda_py311

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cuda.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config rbe)
+TFCI_BUILD_PIP_PACKAGE_ARGS=()
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=(--gpus all)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.11
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-gpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
+TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI=
+TFCI_UPLOAD_LIB_URI=
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/continuous_linux_x86_cuda_py39

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cuda.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config rbe)
+TFCI_BUILD_PIP_PACKAGE_ARGS=()
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=(--gpus all)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.9
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-gpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=
+TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI=
+TFCI_UPLOAD_LIB_URI=
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/nightly_libtensorflow_linux_x86_cpu

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cpu.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config sigbuild_remote_cache_push)
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag)
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=()
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_NVIDIA_SMI_ENABLE=
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI="gs://libtensorflow-nightly/latest"
+TFCI_UPLOAD_LIB_URI="gs://libtensorflow-nightly/$(date -I)"
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/envs/nightly_cpu_py310 → ci/official/envs/nightly_linux_x86_cpu_py310

@@ -12,7 +12,7 @@ TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
 TFCI_INDEX_HTML_ENABLE=1
 TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
-TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_NVIDIA_SMI_ENABLE=
 TFCI_UPLOAD_LIB_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_URI="gs://libtensorflow-nightly/latest"

ci/official/envs/nightly_cpu_py311 → ci/official/envs/nightly_linux_x86_cpu_py311

@@ -12,7 +12,7 @@ TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
 TFCI_INDEX_HTML_ENABLE=1
 TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
-TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_NVIDIA_SMI_ENABLE=
 TFCI_UPLOAD_LIB_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_URI="gs://libtensorflow-nightly/latest"

ci/official/envs/nightly_cpu_py39 → ci/official/envs/nightly_linux_x86_cpu_py39

@@ -12,7 +12,7 @@ TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
 TFCI_INDEX_HTML_ENABLE=1
 TFCI_LIB_SUFFIX="-cpu-linux-x86_64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
-TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_NVIDIA_SMI_ENABLE=
 TFCI_UPLOAD_LIB_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_ENABLE=
 TFCI_UPLOAD_LIB_LATEST_URI="gs://libtensorflow-nightly/latest"

ci/official/envs/nightly_linux_x86_cuda_py310

@@ -0,0 +1,23 @@
+#TFCI_UPLOAD_LIB_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+#TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-release-packages/$RELEASE_VERSION/$KOKORO_GIT_COMMIT_tensorflow"
+TFCI_BAZEL_BAZELRC_ARGS=(--bazelrc ./ci/official/bazelrcs/cuda.bazelrc)
+TFCI_BAZEL_COMMON_ARGS=(--config sigbuild_remote_cache_push)
+TFCI_BUILD_PIP_PACKAGE_ARGS=(--nightly_flag)
+TFCI_COPYBARA_ENABLE=0
+TFCI_DOCKER_ENABLE=1
+TFCI_DOCKER_GPU_ARGS=(--gpus all)
+TFCI_DOCKER_IMAGE=tensorflow/build:latest-python3.10
+TFCI_DOCKER_PULL_ENABLE=1
+TFCI_GIT_DIR=$KOKORO_ARTIFACTS_DIR/github/tensorflow
+TFCI_INDEX_HTML_ENABLE=1
+TFCI_LIB_SUFFIX="-gpu-linux-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_NVIDIA_SMI_ENABLE=1
+TFCI_UPLOAD_LIB_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_ENABLE=
+TFCI_UPLOAD_LIB_LATEST_URI="gs://libtensorflow-nightly/latest"
+TFCI_UPLOAD_LIB_URI="gs://libtensorflow-nightly/$(date -I)"
+TFCI_UPLOAD_WHL_GCS_ENABLE=
+TFCI_UPLOAD_WHL_GCS_URI=
+TFCI_UPLOAD_WHL_PYPI_ARGS=( --config-file "$KOKORO_KEYSTORE_DIR/73361_tensorflow_pypirc_using_global_api_token" --repository testpypi )
+TFCI_UPLOAD_WHL_PYPI_ENABLE=

ci/official/requirements_updater/BUILD.bazel

@@ -13,18 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 
-load("@python//3.8:defs.bzl", compile_pip_requirements_3_8 = "compile_pip_requirements")
 load("@python//3.9:defs.bzl", compile_pip_requirements_3_9 = "compile_pip_requirements")
 load("@python//3.10:defs.bzl", compile_pip_requirements_3_10 = "compile_pip_requirements")
 load("@python//3.11:defs.bzl", compile_pip_requirements_3_11 = "compile_pip_requirements")
 
-compile_pip_requirements_3_8(
-    name = "requirements_3_8",
-    extra_args = ["--allow-unsafe"],
-    requirements_in = "requirements.in",
-    requirements_txt = "requirements_lock_3_8.txt",
-)
-
 compile_pip_requirements_3_9(
     name = "requirements_3_9",
     extra_args = ["--allow-unsafe"],
@@ -45,3 +37,39 @@ compile_pip_requirements_3_11(
     requirements_in = "requirements.in",
     requirements_txt = "requirements_lock_3_11.txt",
 )
+
+compile_pip_requirements_3_9(
+    name = "requirements_3_9_release",
+    extra_args = [
+        "--allow-unsafe",
+        "-P keras-nightly",
+        "-P tb-nightly",
+        "-P tf-estimator-nightly",
+    ],
+    requirements_in = "requirements.in",
+    requirements_txt = "requirements_lock_3_9.txt",
+)
+
+compile_pip_requirements_3_10(
+    name = "requirements_3_10_release",
+    extra_args = [
+        "--allow-unsafe",
+        "-P keras-nightly",
+        "-P tb-nightly",
+        "-P tf-estimator-nightly",
+    ],
+    requirements_in = "requirements.in",
+    requirements_txt = "requirements_lock_3_10.txt",
+)
+
+compile_pip_requirements_3_11(
+    name = "requirements_3_11_release",
+    extra_args = [
+        "--allow-unsafe",
+        "-P keras-nightly",
+        "-P tb-nightly",
+        "-P tf-estimator-nightly",
+    ],
+    requirements_in = "requirements.in",
+    requirements_txt = "requirements_lock_3_11.txt",
+)

ci/official/requirements_updater/WORKSPACE

@@ -28,7 +28,6 @@ python_register_multi_toolchains(
     default_version = default_python_version,
     ignore_root_user_error = True,
     python_versions = [
-        "3.8",
         "3.9",
         "3.10",
         "3.11",

ci/official/requirements_updater/release_updater.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+# 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.
+# ==============================================================================
+
+# script to run pip-compile for keras, tensorboard, estimator deps.
+# if there is a change in requirements.in then all lock files will be updated
+# accordingly.
+
+mv BUILD.bazel BUILD
+
+SUPPORTED_VERSIONS=("3_9" "3_10" "3_11")
+
+for VERSION in "${SUPPORTED_VERSIONS[@]}"
+do
+  cp ../../../requirements_lock_"$VERSION".txt "requirements_lock_"$VERSION".txt"
+  bazel run --experimental_convenience_symlinks=ignore //:requirements_"$VERSION"_release.update
+  sed -i '/^#/d' requirements_lock_"$VERSION".txt
+  mv "requirements_lock_"$VERSION".txt" ../../../requirements_lock_"$VERSION".txt
+done
+
+mv BUILD BUILD.bazel

ci/official/requirements_updater/updater.sh

@@ -21,7 +21,7 @@
 
 mv BUILD.bazel BUILD
 
-SUPPORTED_VERSIONS=("3_8" "3_9" "3_10" "3_11")
+SUPPORTED_VERSIONS=("3_9" "3_10" "3_11")
 
 for VERSION in "${SUPPORTED_VERSIONS[@]}"
 do

ci/official/utilities/rename_and_verify_wheels.sh

@@ -19,13 +19,13 @@
 # "manylinux_xyz" into the wheel filename.
 set -euxo pipefail
 
-cd $1
-for wheel in build/*.whl; do
+DIR=$1
+for wheel in $DIR/*.whl; do
   echo "Checking and renaming $wheel..."
   time python3 -m auditwheel repair --plat manylinux2014_x86_64 "$wheel" --wheel-dir build 2>&1 | tee check.txt
 
   # We don't need the original wheel if it was renamed
-  new_wheel=$(grep --extended-regexp --only-matching '/tf/pkg/\S+.whl' check.txt)
+  new_wheel=$(grep --extended-regexp --only-matching '\S+.whl' check.txt | tail -n 1)
   if [[ "$new_wheel" != "$wheel" ]]; then
     rm "$wheel"
     wheel="$new_wheel"