stack: improve error handling and error messages. (#9629)

This PR refactors the `stack` operation implementation by improving its
error message, and returning a status type value.

**Key Changes:**

- `tensor_methods::stack` returns `StatusOr<XLATensorPtr>`
- Improve error messages and error handling
    - Create `CheckStackAtLeastOneTensor` function

Author: ysiraichi

torch_xla/csrc/aten_xla_type.cpp

@@ -14,6 +14,7 @@
 #include <torch/csrc/lazy/core/tensor_util.h>
 #include <torch/csrc/lazy/core/util.h>
 
+#include <iterator>
 #include <mutex>
 #include <optional>
 
@@ -3696,12 +3697,16 @@ at::Tensor XLANativeFunctions::squeeze_copy(const at::Tensor& self,
 at::Tensor XLANativeFunctions::stack(at::TensorList tensors, int64_t dim) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
   at::ScalarType result_type = at::native::result_type(tensors);
-  std::vector<at::Tensor> c_tensors(tensors.size());
-  std::transform(tensors.begin(), tensors.end(), c_tensors.begin(),
-                 [=](const at::Tensor& t) { return t.to(result_type); });
-  XLA_ASSIGN_OR_THROW(std::vector<absl_nonnull XLATensorPtr> xla_c_tensors,
-                      bridge::GetXlaTensors(c_tensors));
-  return bridge::AtenFromXlaTensor(tensor_methods::stack(xla_c_tensors, dim));
+  std::vector<absl_nonnull XLATensorPtr> xla_tensors;
+  std::transform(tensors.begin(), tensors.end(),
+                 std::back_inserter(xla_tensors), [=](const at::Tensor& t) {
+                   XLA_ASSIGN_OR_THROW(absl_nonnull XLATensorPtr xla_t,
+                                       bridge::GetXlaTensor(t.to(result_type)));
+                   return xla_t;
+                 });
+  XLA_ASSIGN_OR_THROW(absl_nonnull XLATensorPtr output,
+                      tensor_methods::stack(xla_tensors, dim));
+  return bridge::AtenFromXlaTensor(std::move(output));
 }
 
 at::Tensor XLANativeFunctions::std(const at::Tensor& self, bool unbiased) {

torch_xla/csrc/ops/index_ops.cpp

@@ -339,8 +339,8 @@ XLATensorPtr IndexByTensors(const XLATensorPtr& base,
       canonical_indices.front()->shape().get().dimensions_size();
   // Stack the indices to allow the whole multi-indexing to be dispatched with a
   // single gather.
-  XLATensorPtr indices_nd =
-      tensor_methods::stack(canonical_indices, indices_rank);
+  XLA_ASSIGN_OR_THROW(absl_nonnull XLATensorPtr indices_nd,
+                      tensor_methods::stack(canonical_indices, indices_rank));
   return XLATensor::Create(
       torch_xla::MakeNode<IndexGet>(base->GetIrValue(),
                                     indices_nd->GetIrValue(), start_dim),
@@ -356,11 +356,11 @@ torch::lazy::Value IndexPutByTensors(
   }
   auto canonical_indices = WrapIndicesOnce(base, indices, start_dim);
   int64_t indices_rank =
-      canonical_indices.front()->shape().get().dimensions_size();
+      canonical_indices.front()->shape().get().dimensions().size();
   // Stack the indices to allow the whole multi-indexing to be dispatched with a
   // single scatter.
-  XLATensorPtr indices_nd =
-      tensor_methods::stack(canonical_indices, indices_rank);
+  XLA_ASSIGN_OR_THROW(absl_nonnull XLATensorPtr indices_nd,
+                      tensor_methods::stack(canonical_indices, indices_rank));
   return torch_xla::MakeNode<Permute>(
       torch_xla::MakeNode<IndexPut>(base->GetIrValue(),
                                     indices_nd->GetIrValue(), start_dim,

torch_xla/csrc/tensor_methods.cpp

@@ -538,6 +538,15 @@ absl::Status CheckRollDimsAndShiftsAreCompatible(
   return absl::OkStatus();
 }
 
+absl::Status CheckStackAtLeastOneTensor(
+    absl::Span<const absl_nonnull XLATensorPtr> tensors) {
+  if (tensors.size() == 0) {
+    return XLA_ERROR_WITH_LOCATION(
+        absl::InvalidArgumentError("stack(): expected at least one tensor."));
+  }
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 //////////////////////////////////////////////////////////////////////////////
@@ -3422,14 +3431,18 @@ XLATensorPtr squeeze(const XLATensorPtr& input, std::vector<int64_t> dims) {
   return view(input, output_dimensions);
 }
 
-XLATensorPtr stack(absl::Span<const XLATensorPtr> tensors, int64_t dim) {
-  XLA_CHECK_GT(tensors.size(), 0);
+absl::StatusOr<absl_nonnull XLATensorPtr> stack(
+    absl::Span<const absl_nonnull XLATensorPtr> tensors, int64_t dim) {
+  XLA_RETURN_IF_ERROR(CheckStackAtLeastOneTensor(tensors));
+
   std::vector<torch::lazy::Value> values;
-  for (auto& tensor : tensors) {
-    values.push_back(tensor->GetIrValue());
-  }
+  std::transform(
+      tensors.begin(), tensors.end(), std::back_inserter(values),
+      [](const absl_nonnull XLATensorPtr t) { return t->GetIrValue(); });
+
   int64_t canonical_dim = torch::lazy::GetCanonicalDimensionIndex(
-      dim, tensors.front()->shape().get().dimensions_size() + 1);
+      dim, tensors.front()->shape().get().dimensions().size() + 1);
+
   return tensors[0]->CreateFrom(
       torch_xla::MakeNode<Stack>(values, canonical_dim));
 }

torch_xla/csrc/tensor_methods.h

@@ -926,7 +926,8 @@ XLATensorPtr squeeze(const XLATensorPtr& input, std::vector<int64_t> dims);
 void squeeze_(XLATensorPtr& input);
 void squeeze_(XLATensorPtr& input, int64_t dim);
 
-XLATensorPtr stack(absl::Span<const XLATensorPtr> tensors, int64_t dim);
+absl::StatusOr<absl_nonnull XLATensorPtr> stack(
+    absl::Span<const absl_nonnull XLATensorPtr> tensors, int64_t dim);
 
 XLATensorPtr std(const XLATensorPtr& input, std::vector<int64_t> dimensions,
                  bool keep_reduced_dimensions, double correction);

torch_xla/csrc/tensor_ops.cpp

@@ -62,7 +62,9 @@ XLATensorPtr Cross(const XLATensorPtr& input, const XLATensorPtr& other,
   XLATensorPtr s3 = tensor_methods::sub(tensor_methods::mul(u1, v2),
                                         tensor_methods::mul(u2, v1), one);
   // Stack the terms into one result tensor.
-  return tensor_methods::stack({s1, s2, s3}, canonical_dim);
+  XLA_ASSIGN_OR_THROW(absl_nonnull XLATensorPtr output,
+                      tensor_methods::stack({s1, s2, s3}, canonical_dim));
+  return output;
 }
 
 XLATensorPtr MakeMatrixWithDiagonal(const XLATensorPtr& input,