From 0d5dc69ecd62ab37d33fa035919d12522f187d72 Mon Sep 17 00:00:00 2001
From: Yukio Siraichi <yukio.siraichi@gmail.com>
Date: Thu, 7 Aug 2025 16:12:10 -0300
Subject: [PATCH] Refactor status error message builder.

---
 torch_xla/csrc/init_python_bindings.cpp | 13 ++++++++-----
 torch_xla/csrc/status.cpp               | 26 ++++++++++---------------
 torch_xla/csrc/status.h                 |  5 +----
 3 files changed, 19 insertions(+), 25 deletions(-)

diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
index 1d9a4b3c3eb..a2b799a5f0e 100644
--- a/torch_xla/csrc/init_python_bindings.cpp
+++ b/torch_xla/csrc/init_python_bindings.cpp
@@ -3378,13 +3378,16 @@ void InitXlaModuleBindings(py::module m) {
              absl::StatusOr<std::vector<absl_nonnull XLATensorPtr>>
                  xtensors_status = bridge::GetXlaTensors(tensors);
              ABSL_CHECK(xtensors_status.ok())
-                 << "_get_graph_hash(): error retrieving the XLA tensors from "
-                 << "the given tensor arguments. "
+                 << "\n\n"
+                 << "Internal Error:\n"
+                 << "    _get_graph_hash(): error retrieving the XLA tensors "
+                    "from the given tensor arguments. "
                  << "This is a bug! Please, open an issue in the PyTorch/XLA "
                  << "GitHub repository: https://github.com/pytorch/xla"
-                 << std::endl
-                 << "Status Error: "
-                 << BuildStatusErrorMessage(xtensors_status.status());
+                 << "\n\n"
+                 << "Status Error:\n"
+                 << "    " << BuildStatusErrorMessage(xtensors_status.status())
+                 << "\n";
              std::vector<absl_nonnull XLATensorPtr> xtensors =
                  xtensors_status.value();
              torch::lazy::hash_t hash =
diff --git a/torch_xla/csrc/status.cpp b/torch_xla/csrc/status.cpp
index a4e4eddc0cc..2e1c7002e89 100644
--- a/torch_xla/csrc/status.cpp
+++ b/torch_xla/csrc/status.cpp
@@ -103,29 +103,23 @@ static std::string GetFormattedStatusPropagationTrace(
   auto status_propagation_trace = GetStatusPropagationTraceOrEmpty(status);
   return status_propagation_trace.empty()
              ? ""
-             : absl::StrCat("\nStatus Propagation Trace:",
-                            status_propagation_trace.Flatten(), "\n");
-}
-
-// Get the status message followed by a line break, if we are printing the
-// C++ stacktraces.
-//
-// This is needed so we have a blank line in between the status message and
-// the dumped C++ traces (either the status propagation one, or the C++
-// stacktrace).
-static std::string MaybeGetMessageWithLineBreak(const absl::Status& status) {
-  return torch::get_cpp_stacktraces_enabled()
-             ? absl::StrCat(status.message(), "\n")
-             : std::string(status.message());
+             : absl::StrCat("\n\nStatus Propagation Trace:",
+                            status_propagation_trace.Flatten());
 }
 
 std::string BuildStatusErrorMessage(const absl::Status& status) {
-  return absl::StrCat(MaybeGetMessageWithLineBreak(status),
+  return absl::StrCat(status.message(),
                       GetFormattedStatusPropagationTrace(status));
 }
 
+// Return a line break if torch::get_cpp_stacktraces_enabled() is true.
+static std::string LineBreakIfCppStacktracesEnabled() {
+  return torch::get_cpp_stacktraces_enabled() ? "\n" : "";
+}
+
 void MaybeThrow(const absl::Status& status) {
-  TORCH_CHECK(status.ok(), BuildStatusErrorMessage(status));
+  TORCH_CHECK(status.ok(), absl::StrCat(BuildStatusErrorMessage(status),
+                                        LineBreakIfCppStacktracesEnabled()));
 }
 
 void GetValueOrThrow(const absl::Status& status) { MaybeThrow(status); }
diff --git a/torch_xla/csrc/status.h b/torch_xla/csrc/status.h
index c922e2f511f..b2d508076a3 100644
--- a/torch_xla/csrc/status.h
+++ b/torch_xla/csrc/status.h
@@ -179,10 +179,7 @@ absl::Status MaybeWithNewMessage(const absl::Status& status, const char* file,
 // If `TORCH_SHOW_CPP_STACKTRACES` is enabled, returns the concatenation of
 // `status.message()` with its inner status propagation trace.
 //
-// TODO(ysiraichi): this call does not append the C++ stacktrace, which,
-// ideally, should. It can be done by not using `TORCH_CHECK()` macro directly
-// in `MaybeThrow()`, but using PyTorch `c10::get_lazy_backtrace()`
-// (at c10/util/Backtrace.h).
+// It doesn't add a trailing line break.
 std::string BuildStatusErrorMessage(const absl::Status& status);
 
 // Maybe throws an exception if `status` has a non-ok code.