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

…ync-230515

Develop upstream sync 230515

.bazelrc

@@ -446,7 +446,7 @@ build:rbe --bes_backend=buildeventservice.googleapis.com
 build:rbe --bes_results_url="https://source.cloud.google.com/results/invocations"
 build:rbe --bes_timeout=600s
 build:rbe --define=EXECUTOR=remote
-build:rbe --flaky_test_attempts=3
+build:rbe --flaky_test_attempts=2
 build:rbe --jobs=800
 build:rbe --remote_executor=grpcs://remotebuildexecution.googleapis.com
 build:rbe --remote_timeout=3600
@@ -627,7 +627,7 @@ try-import %workspace%/.bazelrc.user
 
 # Here are bazelrc configs for release builds
 build:release_base --config=v2
-test:release_base --flaky_test_attempts=3
+test:release_base --flaky_test_attempts=2
 test:release_base --test_size_filters=small,medium
 
 build:release_cpu_linux --config=release_base

.github/workflows/arm-ci-extended.yml

@@ -28,7 +28,7 @@ jobs:
     runs-on: [self-hosted, linux, ARM64]
     strategy:
       matrix:
-        pyver: ['3.10']
+        pyver: ['3.8', '3.9', '3.10', '3.11']
     steps:
       - name: Stop old running containers (if any)
         shell: bash

.github/workflows/update-rbe.yml

@@ -92,6 +92,18 @@ jobs:
         map sigbuild-r2.13-clang-python3.9 2.13-python3.9
         map sigbuild-r2.13-clang-python3.10 2.13-python3.10
         map sigbuild-r2.13-clang-python3.11 2.13-python3.11
+        # TF 2.14
+        map sigbuild-r2.14 2.14-python3.9
+        map sigbuild-r2.14-python3.8 2.14-python3.8
+        map sigbuild-r2.14-python3.9 2.14-python3.9
+        map sigbuild-r2.14-python3.10 2.14-python3.10
+        map sigbuild-r2.14-python3.11 2.14-python3.11
+        # TF 2.14 + Clang (containers are the same, but env vars in configs.bzl are different)
+        map sigbuild-r2.14-clang 2.14-python3.9
+        map sigbuild-r2.14-clang-python3.8 2.14-python3.8
+        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
     - name: Create Pull Request with changes
       uses: peter-evans/create-pull-request@2b011faafdcbc9ceb11414d64d0573f37c774b04 # v4.2.3
       with:

README.md

@@ -2,7 +2,7 @@
   <img src="https://www.tensorflow.org/images/tf_logo_horizontal.png">
 </div>
 
-[![Python](https://img.shields.io/pypi/pyversions/tensorflow.svg?style=plastic)](https://badge.fury.io/py/tensorflow)
+[![Python](https://img.shields.io/pypi/pyversions/tensorflow.svg)](https://badge.fury.io/py/tensorflow)
 [![PyPI](https://badge.fury.io/py/tensorflow.svg)](https://badge.fury.io/py/tensorflow)
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.4724125.svg)](https://doi.org/10.5281/zenodo.4724125)
 [![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/1486/badge)](https://bestpractices.coreinfrastructure.org/projects/1486)
@@ -11,6 +11,8 @@
 [![Fuzzing Status](https://oss-fuzz-build-logs.storage.googleapis.com/badges/tensorflow-py.svg)](https://bugs.chromium.org/p/oss-fuzz/issues/list?sort=-opened&can=1&q=proj:tensorflow-py)
 [![OSSRank](https://shields.io/endpoint?url=https://ossrank.com/shield/44)](https://ossrank.com/p/44)
 [![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-v1.4%20adopted-ff69b4.svg)](CODE_OF_CONDUCT.md)
+[![TF Official Continuous](https://tensorflow.github.io/build/TF%20Official%20Continuous.svg)](https://tensorflow.github.io/build#TF%20Official%20Continuous)
+[![TF Official Nightly](https://tensorflow.github.io/build/TF%20Official%20Nightly.svg)](https://tensorflow.github.io/build#TF%20Official%20Nightly)
 
 **`Documentation`** |
 ------------------- |

RELEASE.md

@@ -34,6 +34,18 @@
 * <IF A CHANGE CLOSES A GITHUB ISSUE, IT SHOULD BE DOCUMENTED HERE>
 * <NOTES SHOULD BE GROUPED PER AREA>
 
+* `tf.config.experimental.enable_tensor_float_32_execution`
+    * Disabling TensorFloat-32 execution now causes TPUs to use float32
+      precision for float32 matmuls and other ops. TPUs have always used
+      bfloat16 precision for certain ops, like matmul, when such ops had float32
+      inputs. Now, disabling TensorFloat-32 by calling
+      `tf.config.experimental.enable_tensor_float_32_execution(False)` will
+      cause TPUs to use float32 precision for such ops instead of bfloat16.
+
+*  `tf.experimental.dtensor`
+    * API changes for Relayout. Added a new API, `dtensor.relayout_like`, for 
+      relayouting a tensor according to the layout of another tensor. 
+
 # Thanks to our Contributors
 
 This release contains contributions from many people at Google, as well as:
@@ -243,6 +255,8 @@ This release contains contributions from many people at Google, as well as:
 
 *   `tf.lite`:
     *   Add UINT32 support to tfl.pack
+    *   Add INT64 support to tfl.range
+    *   Add UINT32 support to tfl.concatenation
 
 ## Thanks to our Contributors
 

configure.py

@@ -964,7 +964,7 @@ def set_other_cuda_vars(environ_cp):
 
 def system_specific_test_config(environ_cp):
   """Add default build and test flags required for TF tests to bazelrc."""
-  write_to_bazelrc('test --flaky_test_attempts=3')
+  write_to_bazelrc('test --flaky_test_attempts=2')
   write_to_bazelrc('test --test_size_filters=small,medium')
 
   # Each instance of --test_tag_filters or --build_tag_filters overrides all

tensorflow/c/BUILD

@@ -968,7 +968,7 @@ tf_cc_test(
 
 tf_cc_test(
     name = "c_api_function_test",
-    size = "small",
+    size = "medium",
     srcs = ["c_api_function_test.cc"],
     deps = [
         ":c_api",
@@ -985,7 +985,7 @@ tf_cc_test(
 
 tf_cc_test(
     name = "while_loop_test",
-    size = "small",
+    size = "medium",
     srcs = ["while_loop_test.cc"],
     deps = [
         ":c_api",

tensorflow/c/c_api_internal.h

@@ -142,7 +142,7 @@ struct TF_ImportGraphDefOptions {
 
   // Backing memory for TensorId fields in opts.
   // TODO(skyewm): it'd be better if ImportGraphDefOptions owned this.
-  std::list<tensorflow::string> tensor_id_data;
+  std::vector<tensorflow::string> tensor_id_data;
 };
 
 struct TF_ImportGraphDefResults {
@@ -152,7 +152,7 @@ struct TF_ImportGraphDefResults {
   std::vector<int> missing_unused_key_indexes;
 
   // Backing memory for missing_unused_key_names values.
-  std::list<tensorflow::string> missing_unused_key_names_data;
+  std::vector<tensorflow::string> missing_unused_key_names_data;
 };
 
 struct TF_DeviceList {

tensorflow/c/c_api_macros.h

@@ -26,7 +26,12 @@ limitations under the License.
 #define TF_CAPI_EXPORT __declspec(dllimport)
 #endif  // TF_COMPILE_LIBRARY
 #else
+#ifdef TF_CAPI_WEAK
+#define TF_CAPI_EXPORT \
+  __attribute__((visibility("default"))) __attribute((weak))
+#else
 #define TF_CAPI_EXPORT __attribute__((visibility("default")))
+#endif  // TF_CAPI_WEAK
 #endif  // _WIN32
 #endif  // SWIG
 

tensorflow/c/eager/BUILD

@@ -725,7 +725,7 @@ tf_cuda_cc_test(
 
 tf_cuda_cc_test(
     name = "c_api_remote_function_test",
-    size = "small",
+    size = "medium",
     srcs = [
         "c_api_remote_function_test.cc",
     ],

tensorflow/c/experimental/next_pluggable_device/BUILD

@@ -32,6 +32,8 @@ cc_library(
         "//tensorflow/tsl/distributed_runtime/coordination:coordination_service_agent",
         "//tensorflow/tsl/platform:errors",
         "//tensorflow/tsl/platform:statusor",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
     ],
 )
 

tensorflow/c/experimental/next_pluggable_device/c_api.cc

@@ -21,6 +21,8 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/c/kernels_experimental.h"
 #include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/c/tf_status_internal.h"
@@ -110,9 +112,9 @@ TF_VariableInfo* TF_CreateVariableInfoFromContext(TF_OpKernelContext* ctx,
   const tensorflow::Tensor& arg_tensor = cc_ctx->input(index);
   tsl::Status cc_status;
   if (arg_tensor.dtype() != tensorflow::DT_RESOURCE) {
-    cc_status = tsl::errors::InvalidArgument(
-        "Trying to obtain resource handle from Input[", index,
-        "], which is not type DT_RESOURCE.");
+    cc_status = absl::InvalidArgumentError(
+        absl::StrCat("Trying to obtain resource handle from Input[", index,
+                     "], which is not type DT_RESOURCE."));
     tsl::Set_TF_Status_from_Status(status, cc_status);
     return nullptr;
   }
@@ -140,12 +142,12 @@ void TF_AllocateTempForVariableInfo(TF_OpKernelContext* ctx,
   auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
   tsl::Status cc_status;
   if (var_info == nullptr) {
-    cc_status = tsl::errors::InvalidArgument("TF_VariableInfo is NULL.");
+    cc_status = absl::InvalidArgumentError("TF_VariableInfo is NULL.");
     tsl::Set_TF_Status_from_Status(status, cc_status);
     return;
   }
   if (var_info->var_info.var() == nullptr) {
-    cc_status = tsl::errors::InvalidArgument(
+    cc_status = absl::InvalidArgumentError(
         "VariableInfo does not track a resource variable.");
     tsl::Set_TF_Status_from_Status(status, cc_status);
     return;
@@ -161,12 +163,12 @@ TF_Tensor* TF_GetTensorFromVariableInfo(TF_VariableInfo* var_info,
                                         TF_Status* status) {
   tsl::Status cc_status;
   if (var_info == nullptr) {
-    cc_status = tsl::errors::InvalidArgument("TF_VariableInfo is NULL.");
+    cc_status = absl::InvalidArgumentError("TF_VariableInfo is NULL.");
     tsl::Set_TF_Status_from_Status(status, cc_status);
     return nullptr;
   }
   if (var_info->var_info.var() == nullptr) {
-    cc_status = tsl::errors::InvalidArgument(
+    cc_status = absl::InvalidArgumentError(
         "VariableInfo does not track a resource variable.");
     tsl::Set_TF_Status_from_Status(status, cc_status);
     return nullptr;
@@ -263,8 +265,9 @@ PJRT_Client* TF_GetPjRtCClient(const char* device_type, TF_Status* status) {
       tensorflow::down_cast<xla::PjRtCApiClient*>(*pjrt_client);
   if (pjrt_c_api_client == nullptr) {
     tensorflow::Set_TF_Status_from_Status(
-        status, tsl::errors::Internal("PjRtClient for ", device_type,
-                                      " is not type PjRtCApiClient"));
+        status,
+        absl::InternalError(absl::StrCat("PjRtClient for ", device_type,
+                                         " is not type PjRtCApiClient")));
     return nullptr;
   }
   TF_SetStatus(status, TF_OK, "");
@@ -282,16 +285,15 @@ PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor, TF_Status* status) {
       tensorflow::AsyncValueTensor::FromTensor(&tensor);
   if (av_tensor == nullptr || av_tensor->GetBuffer() == nullptr) {
     tensorflow::Set_TF_Status_from_Status(
-        status,
-        tsl::errors::Internal("Input tensor does not have PjRtBuffer."));
+        status, absl::InternalError("Input tensor does not have PjRtBuffer."));
     return nullptr;
   }
   auto* c_api_buffer =
       tensorflow::down_cast<xla::PjRtCApiBuffer*>(av_tensor->GetBuffer().get());
   if (c_api_buffer == nullptr) {
     tensorflow::Set_TF_Status_from_Status(
         status,
-        tsl::errors::Internal(
+        absl::InternalError(
             "The PjRtBuffer in the tensor is not type PjRtCApiBuffer."));
     return nullptr;
   }
@@ -317,16 +319,17 @@ void TF_CreatePjRtBuffer(TF_Tensor* c_tensor, PJRT_Buffer* c_buffer,
       tensorflow::down_cast<xla::PjRtCApiClient*>(*pjrt_client);
   if (pjrt_c_api_client == nullptr) {
     tensorflow::Set_TF_Status_from_Status(
-        status, tsl::errors::Internal("PjRtClient for ", device_type,
-                                      " is not type PjRtCApiClient"));
+        status,
+        absl::InternalError(absl::StrCat("PjRtClient for ", device_type,
+                                         " is not type PjRtCApiClient")));
     return;
   }
   tensorflow::AsyncValueTensor* av_tensor =
       tensorflow::AsyncValueTensor::FromTensor(&tensor);
   if (av_tensor == nullptr) {
     tensorflow::Set_TF_Status_from_Status(
         status,
-        tsl::errors::Internal(
+        absl::InternalError(
             "The tensor to set PjRtBuffer is not an AsyncValueTensor."));
     return;
   }

tensorflow/c/experimental/saved_model/core/revived_types/BUILD

@@ -28,6 +28,8 @@ cc_library(
         "//tensorflow/cc/saved_model:constants",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
     ],
 )
 
@@ -99,6 +101,8 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/lib/llvm_rtti",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
     ],
 )

tensorflow/c/experimental/saved_model/core/revived_types/asset.cc

@@ -17,6 +17,8 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/c/eager/immediate_execution_context.h"
 #include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
 #include "tensorflow/c/tensor_interface.h"
@@ -37,8 +39,8 @@ Status Asset::Create(ImmediateExecutionContext* ctx,
       io::JoinPath(saved_model_dir, kSavedModelAssetsDirectory, asset_filename);
   AbstractTensorPtr tensor(ctx->CreateStringScalar(abs_path));
   if (tensor.get() == nullptr) {
-    return errors::Internal(
-        "Failed to create scalar string tensor for Asset at path ", abs_path);
+    return absl::InternalError(absl::StrCat(
+        "Failed to create scalar string tensor for Asset at path ", abs_path));
   }
 
   ImmediateTensorHandlePtr handle(ctx->CreateLocalHandle(tensor.get()));