Improve error message for shape promotion on lowering. (pytorch#9486)

Author: ysiraichi

test/test_operations.py

@@ -2445,6 +2445,19 @@ def test_isneginf_no_fallback(self):
     t = t.to(torch.float16)
     self._test_no_fallback(torch.isneginf, (t,))
 
+  def test_add_broadcast_error(self):
+    a = torch.rand(2, 2, 4, 4, device="xla")
+    b = torch.rand(2, 2, device="xla")
+
+    expected_regex = (
+        r"Shapes are not compatible for broadcasting: f32\[2,2,4,4\] vs. f32\[2,2\]. "
+        r"Expected dimension 2 of shape f32\[2,2,4,4\] \(4\) to match dimension "
+        r"0 of shape f32\[2,2\] \(2\). .*")
+
+    with self.assertRaisesRegex(RuntimeError, expected_regex):
+      torch.add(a, b)
+      torch_xla.sync()
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/data_ops.cpp

@@ -17,6 +17,7 @@
 #include "torch_xla/csrc/runtime/sys_util.h"
 #include "torch_xla/csrc/runtime/util.h"
 #include "torch_xla/csrc/shape_helper.h"
+#include "torch_xla/csrc/status.h"
 #include "torch_xla/csrc/tensor_util.h"
 #include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/lib/slicing.h"
@@ -196,7 +197,8 @@ xla::XlaOp BuildMaskedFillScalar(xla::XlaOp input, xla::XlaOp mask,
   const xla::Shape& mask_shape = ShapeHelper::ShapeOfXlaOp(mask);
 
   if (!xla::ShapeUtil::Compatible(input_shape, mask_shape)) {
-    xla::Shape shape = XlaHelpers::GetPromotedShape(input_shape, mask_shape);
+    xla::Shape shape =
+        GetValueOrThrow(XlaHelpers::GetPromotedShape(input_shape, mask_shape));
     input = BuildExpand(input, shape.dimensions());
     mask = BuildExpand(mask, shape.dimensions());
   }

torch_xla/csrc/helpers.cpp

@@ -582,8 +582,8 @@ void ExtractDimensionSizesAndDynamicDimensionsFromShape(
 
 }  // namespace
 
-xla::Shape XlaHelpers::GetPromotedShape(const xla::Shape& shape1,
-                                        const xla::Shape& shape2) {
+absl::StatusOr<xla::Shape> XlaHelpers::GetPromotedShape(
+    const xla::Shape& shape1, const xla::Shape& shape2) {
   std::vector<int64_t> dimensions;
   std::vector<bool> dynamic_dimensions;
 
@@ -606,20 +606,33 @@ xla::Shape XlaHelpers::GetPromotedShape(const xla::Shape& shape1,
   size_t min_size =
       std::min(shape1.dimensions().size(), shape2.dimensions().size());
   for (size_t i = 0; i < min_size; i++) {
-    int64_t dim1 =
-        shape1.dimensions()[shape1.dimensions().size() - min_size + i];
+    int64_t dim_index1 = shape1.dimensions().size() - min_size + i;
+    int64_t dim_index2 = shape2.dimensions().size() - min_size + i;
+    int64_t dim1 = shape1.dimensions()[dim_index1];
+    int64_t dim2 = shape2.dimensions()[dim_index2];
+
     int64_t dynamic_dim1 =
         shape1.dynamic_dimensions()[shape1.dynamic_dimensions().size() -
                                     min_size + i];
-    int64_t dim2 =
-        shape2.dimensions()[shape2.dimensions().size() - min_size + i];
     int64_t dynamic_dim2 =
         shape2.dynamic_dimensions()[shape2.dynamic_dimensions().size() -
                                     min_size + i];
 
-    XLA_CHECK(dim1 == dim2 || dim1 == 1 || dim2 == 1 ||
-              dim1 == xla::Shape::kUnboundedSize ||
-              dim2 == xla::Shape::kUnboundedSize);
+    if (dim1 != dim2 && dim1 != 1 && dim2 != 1 &&
+        dim1 != xla::Shape::kUnboundedSize &&
+        dim2 != xla::Shape::kUnboundedSize) {
+      auto shape_str1 = shape1.ToString();
+      auto shape_str2 = shape2.ToString();
+      auto message = absl::StrCat(
+          "Shapes are not compatible for broadcasting: ", shape_str1, " vs. ",
+          shape_str2, ". Expected dimension ", dim_index1, " of shape ",
+          shape_str1, " (", dim1, ") ", "to match dimension ", dim_index2,
+          " of shape ", shape_str2, " (", dim2, "). ",
+          "Either that or that any of them is either 1 or unbounded. ",
+          "Try reshaping one of the tensors to match the "
+          "other.");
+      return XLA_ERROR_WITH_LOCATION(absl::InvalidArgumentError(message));
+    }
 
     // TODO: Consider replacing the broadcasting logic below with
     // 'xla::ShapeInference::InferDegenerateDimensionBroadcastShape' resuing the
@@ -684,7 +697,7 @@ std::vector<int64_t> XlaHelpers::getBroadcastDimensions(xla::XlaOp op1,
 xla::Shape XlaHelpers::GetPromotedBinaryOpShape(const xla::Shape& shape1,
                                                 const xla::Shape& shape2) {
   if (!shape1.is_dynamic() && !shape2.is_dynamic()) {
-    auto promoted_shape = GetPromotedShape(shape1, shape2);
+    auto promoted_shape = GetValueOrThrow(GetPromotedShape(shape1, shape2));
     return xla::ShapeUtil::MakeShape(
         PromoteType(shape1.element_type(), shape2.element_type()),
         promoted_shape.dimensions());
@@ -763,7 +776,7 @@ std::pair<xla::XlaOp, xla::XlaOp> XlaHelpers::PromoteShapes(xla::XlaOp op1,
   const xla::Shape& shape1 = ShapeHelper::ShapeOfXlaOp(op1);
   const xla::Shape& shape2 = ShapeHelper::ShapeOfXlaOp(op2);
 
-  xla::Shape shape = GetPromotedShape(shape1, shape2);
+  xla::Shape shape = GetValueOrThrow(GetPromotedShape(shape1, shape2));
   if (shape1.is_unbounded_dynamic() || shape2.is_unbounded_dynamic()) {
     return ImplicitBroadcastWithUnboundedDynamicShapes(op1, op2, shape);
   }

torch_xla/csrc/helpers.h

@@ -11,6 +11,7 @@
 #include <vector>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/types/optional.h"
 #include "absl/types/span.h"
 #include "torch_xla/csrc/runtime/debug_macros.h"
@@ -302,8 +303,8 @@ class XlaHelpers {
                                                          xla::XlaOp op2);
 
   // Given the two shape 'shape1' and 'shape2', infers the broadcasted shape.
-  static xla::Shape GetPromotedShape(const xla::Shape& shape1,
-                                     const xla::Shape& shape2);
+  static absl::StatusOr<xla::Shape> GetPromotedShape(const xla::Shape& shape1,
+                                                     const xla::Shape& shape2);
 
   static xla::Shape GetPromotedDynamicShape(const xla::Shape& shape1,
                                             const xla::Shape& shape2);

torch_xla/csrc/ops/triangular_solve.cpp

@@ -3,6 +3,7 @@
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/lowering_context.h"
 #include "torch_xla/csrc/shape_helper.h"
+#include "torch_xla/csrc/status.h"
 #include "xla/hlo/builder/xla_builder.h"
 #include "xla/layout_util.h"
 
@@ -32,8 +33,8 @@ std::pair<xla::Shape, xla::Shape> InferTriangularSolveShape(
     return std::pair<xla::Shape, xla::Shape>(rhs_batch_shape, lhs_batch_shape);
   }
   // Obtain the promoted shapes and add back the trailing dimension.
-  xla::Shape rhs_batch_promoted_shape =
-      XlaHelpers::GetPromotedShape(rhs_batch_shape, lhs_batch_shape);
+  xla::Shape rhs_batch_promoted_shape = GetValueOrThrow(
+      XlaHelpers::GetPromotedShape(rhs_batch_shape, lhs_batch_shape));
   xla::Shape lhs_batch_promoted_shape(rhs_batch_promoted_shape);
   rhs_batch_promoted_shape.add_dimensions(nrhs);
   lhs_batch_promoted_shape.add_dimensions(n);

torch_xla/csrc/xla_lower_util.cpp

@@ -63,7 +63,7 @@ ConditionMaskData CreateConditionMaskData(xla::XlaOp condition) {
 xla::XlaOp GetPromotedMask(xla::XlaOp mask, const xla::Shape& input_shape) {
   const xla::Shape& mask_shape = ShapeHelper::ShapeOfXlaOp(mask);
   xla::Shape promoted_mask_shape =
-      XlaHelpers::GetPromotedShape(mask_shape, input_shape);
+      GetValueOrThrow(XlaHelpers::GetPromotedShape(mask_shape, input_shape));
   return XlaHelpers::ImplicitBroadcast(mask, mask_shape, promoted_mask_shape);
 }
 
@@ -543,7 +543,8 @@ std::vector<xla::XlaOp> CreateBroadcastTensors(
   for (const xla::XlaOp operand : operands) {
     const xla::Shape& operand_shape = ShapeHelper::ShapeOfXlaOp(operand);
     operand_shapes.push_back(operand_shape);
-    result_shape = XlaHelpers::GetPromotedShape(result_shape, operand_shape);
+    result_shape = GetValueOrThrow(
+        XlaHelpers::GetPromotedShape(result_shape, operand_shape));
   }
   std::vector<xla::XlaOp> result;
   for (size_t i = 0; i < operands.size(); ++i) {