Extend status propagation to XlaDataToTensors and update callers

Author: ysiraichi

test/cpp/test_xla_sharding.cpp

@@ -29,7 +29,7 @@ bool XlaDataValuesEqual(torch::lazy::BackendDataPtr a,
                         torch::lazy::BackendDataPtr b,
                         at::ScalarType element_type) {
   std::vector<at::Tensor> tensors =
-      XlaDataToTensors({a, b}, {element_type, element_type});
+      GetValueOrThrow(XlaDataToTensors({a, b}, {element_type, element_type}));
   return TensorCompare(tensors[0], tensors[1]);
 }
 }  // namespace

torch_xla/csrc/init_python_bindings.cpp

@@ -2712,7 +2712,7 @@ void InitXlaModuleBindings(py::module m) {
             }
 
             std::vector<at::Tensor> cpu_shards =
-                XlaDataToTensors(WrapXlaData(handles), element_types);
+                GetValueOrThrow(XlaDataToTensors(WrapXlaData(handles), element_types));
             // Populate the resulting vector of shards and device strings
             std::vector<std::vector<std::pair<at::Tensor, std::string>>> result;
             int shards_per_tensor =

torch_xla/csrc/tensor.cpp

@@ -40,6 +40,7 @@
 #include "torch_xla/csrc/runtime/pjrt_computation_client.h"
 #include "torch_xla/csrc/runtime/sys_util.h"
 #include "torch_xla/csrc/runtime/xla_util.h"
+#include "torch_xla/csrc/status.h"
 #include "torch_xla/csrc/tensor_util.h"
 #include "torch_xla/csrc/torch_util.h"
 #include "torch_xla/csrc/xla_graph_executor.h"
@@ -512,7 +513,7 @@ at::Tensor XLATensor::ToTensor(bool detached) {
     // The GetXlaData() call will trigger an ApplyPendingGraph() if an IR
     // XlaNode is available on the tensor.
     std::vector<at::Tensor> tensors =
-        XlaDataToTensors({GetXlaData()}, {dtype()});
+        GetValueOrThrow(XlaDataToTensors({GetXlaData()}, {dtype()}));
     tensor = std::move(tensors.front());
     if (!detached) {
       SetTensorData(tensor);

torch_xla/csrc/tensor_util.cpp

@@ -922,11 +922,11 @@ absl::StatusOr<std::vector<xla::Literal>> ReleaseGilAndTransferData(
   return literals;
 }
 
-std::vector<at::Tensor> XlaDataToTensors(
+absl::StatusOr<std::vector<at::Tensor>> XlaDataToTensors(
     absl::Span<const torch::lazy::BackendDataPtr> xla_data,
     absl::Span<const at::ScalarType> dest_element_type) {
-  std::vector<xla::Literal> literals =
-      GetValueOrThrow(ReleaseGilAndTransferData(xla_data));
+  XLA_ASSIGN_OR_RETURN(std::vector<xla::Literal> literals,
+                       ReleaseGilAndTransferData(xla_data));
   std::vector<at::Tensor> tensors(literals.size());
   absl::BlockingCounter counter(literals.size());
   for (size_t i = 0; i < tensors.size(); ++i) {

torch_xla/csrc/tensor_util.h

@@ -32,7 +32,7 @@ absl::StatusOr<std::vector<xla::Literal>> ReleaseGilAndTransferData(
     absl::Span<const torch::lazy::BackendDataPtr> xla_data);
 
 // TODO LTC @wonjoo - Migrate to upstream after Device -> BackendDevice
-std::vector<at::Tensor> XlaDataToTensors(
+absl::StatusOr<std::vector<at::Tensor>> XlaDataToTensors(
     absl::Span<const torch::lazy::BackendDataPtr> xla_data,
     absl::Span<const at::ScalarType> dest_element_type);
 

torch_xla/csrc/xla_backend_impl.cpp

@@ -10,6 +10,8 @@
 #include "torch_xla/csrc/runtime/computation_client.h"
 #include "torch_xla/csrc/runtime/debug_macros.h"
 #include "torch_xla/csrc/runtime/runtime.h"
+#include "torch_xla/csrc/status.h"
+#include "torch_xla/csrc/tensor_util.h"
 
 namespace at {
 // This function is defined in the codegenerated RegisterDispatchKey.cpp file.
@@ -92,7 +94,7 @@ class XlaBackendImpl : public torch::lazy::BackendImplInterface {
       const torch::lazy::BackendDataPtr data,
       std::optional<at::ScalarType> logical_scalar_type) const override {
     // TODO(JackCaoG): handle the logical_scalar_type == nullptr case
-    return XlaDataToTensors({data}, {*logical_scalar_type})[0];
+    return GetValueOrThrow(XlaDataToTensors({data}, {*logical_scalar_type}))[0];
   }
 
   std::unique_ptr<torch::lazy::LoweringContext> CreateLoweringContext(