Avoid losing payloads at Status recreation

PiperOrigin-RevId: 395088620
Change-Id: Ib17e2942a04351a98533c51a978f3c5227fd68d4

Author: faizan-m

tensorflow/compiler/xla/service/hlo_pass_pipeline.cc

@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/compiler/xla/status_macros.h"
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/compiler/xla/util.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 
 namespace xla {
@@ -130,9 +131,10 @@ Status HloPassPipeline::RunInvariantCheckers(
     if (!changed_status.ok()) {
       VLOG(2) << "Failed invariant check:";
       XLA_VLOG_LINES(2, hlo->ToString());
-      return Status(changed_status.status().code(),
-                    absl::StrCat(changed_status.status().error_message(),
-                                 "\n\nFailed after ", after_pass_name));
+      return tensorflow::errors::CreateWithUpdatedMessage(
+          changed_status.status(),
+          absl::StrCat(changed_status.status().error_message(),
+                       "\n\nFailed after ", after_pass_name));
     }
     TF_RET_CHECK(!changed_status.ValueOrDie())
         << "invariant checkers must not change the graph";

tensorflow/compiler/xla/service/service.cc

@@ -240,9 +240,10 @@ Service::ResolveAndValidateArguments(
   for (size_t i = 0; i < arguments.size(); ++i) {
     auto buffer_status = allocation_tracker_.Resolve(*arguments[i]);
     if (!buffer_status.ok()) {
-      return Status(buffer_status.status().code(),
-                    StrCat(buffer_status.status().error_message(), ", ",
-                           "failed to resolve allocation for parameter ", i));
+      return tensorflow::errors::CreateWithUpdatedMessage(
+          buffer_status.status(),
+          StrCat(buffer_status.status().error_message(), ", ",
+                 "failed to resolve allocation for parameter ", i));
     }
     auto replicated_buffers = buffer_status.ValueOrDie();
     CHECK_EQ(options_.number_of_replicas(), replicated_buffers.size());