div: improve error handling and error messages. (#9549)

Author: ysiraichi

test/test_operations.py

@@ -2486,6 +2486,17 @@ def test_cat_raises_error_on_incompatible_shapes(self):
           "or that either of them was a 1D empty tensor of size (0,).")
       self.assertEqual(str(e), expected_error)
 
+  def test_div_raises_error_on_invalid_rounding_mode(self):
+    a = torch.rand(2, 2, device=torch_xla.device())
+
+    try:
+      torch.div(a, 2, rounding_mode="bad")
+    except RuntimeError as e:
+      expected_error = (
+          "div(): invalid rounding mode `bad`. Expected it to be either "
+          "'trunc', 'floor', or be left unspecified.")
+      self.assertEqual(str(e), expected_error)
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/aten_xla_type.cpp

@@ -1536,8 +1536,9 @@ at::Tensor XLANativeFunctions::div(
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
   at::ScalarType dtype = at::result_type(self, other);
   auto operands = GetBinaryOperands(self, UnwrapNumber(other, dtype));
-  return bridge::AtenFromXlaTensor(tensor_methods::div(
+  auto output = GetValueOrThrow(tensor_methods::div(
       operands.first, operands.second, rounding_mode, dtype));
+  return bridge::AtenFromXlaTensor(std::move(output));
 }
 
 at::Tensor XLANativeFunctions::div(const at::Tensor& self,

torch_xla/csrc/tensor_methods.cpp

@@ -1421,9 +1421,10 @@ XLATensorPtr diagonal(const XLATensorPtr& input, int64_t offset, int64_t dim1,
       input->GetIrValue(), offset, canonical_dim1, canonical_dim2));
 }
 
-XLATensorPtr div(const XLATensorPtr& input, const XLATensorPtr& other,
-                 const std::optional<std::string_view>& rounding_mode,
-                 std::optional<at::ScalarType> logical_element_type) {
+absl::StatusOr<absl_nonnull XLATensorPtr> div(
+    const XLATensorPtr& input, const XLATensorPtr& other,
+    const std::optional<std::string_view>& rounding_mode,
+    std::optional<at::ScalarType> logical_element_type) {
   at::ScalarType scalar_type =
       at::typeMetaToScalarType(c10::get_default_dtype());
   xla::PrimitiveType input_type = input->shape().get().element_type();
@@ -1446,8 +1447,10 @@ XLATensorPtr div(const XLATensorPtr& input, const XLATensorPtr& other,
     } else if (*rounding_mode == "floor") {
       res = torch_xla::MakeNode<Floor>(res);
     } else {
-      XLA_CHECK(false)
-          << "rounding_mode must be one of None, 'trunc', or 'floor'";
+      return XLA_ERROR_WITH_LOCATION(absl::InvalidArgumentError(
+          absl::StrCat("div(): invalid rounding mode `", *rounding_mode,
+                       "`. Expected it to be either 'trunc', 'floor', or be "
+                       "left unspecified.")));
     }
   }
 

torch_xla/csrc/tensor_methods.h

@@ -389,7 +389,7 @@ XLATensorPtr diag(const XLATensorPtr& input, int64_t offset);
 XLATensorPtr diagonal(const XLATensorPtr& input, int64_t offset, int64_t dim1,
                       int64_t dim2);
 
-XLATensorPtr div(
+absl::StatusOr<absl_nonnull XLATensorPtr> div(
     const XLATensorPtr& input, const XLATensorPtr& other,
     const std::optional<std::string_view>& rounding_mode = std::nullopt,
     std::optional<at::ScalarType> logical_element_type = std::nullopt);

torch_xla/csrc/tensor_ops.cpp

@@ -148,9 +148,9 @@ XLATensorPtr SmoothL1LossBackward(const XLATensorPtr& grad_output,
       XLATensorPtr grad_scale = tensor_methods::get_dimensions_size(
           broadcasted_input,
           XlaHelpers::GetAllDimensions(broadcasted_input->shape()));
-      return tensor_methods::mul(
-          tensor_methods::div(elementwise_loss_backward, grad_scale),
-          grad_output);
+      XLATensorPtr div_result = GetValueOrThrow(
+          tensor_methods::div(elementwise_loss_backward, grad_scale));
+      return tensor_methods::mul(div_result, grad_output);
     }
     default:
       XLA_ERROR() << "Invalid reduction type: "
@@ -174,12 +174,12 @@ XLATensorPtr SoftplusBackward(const XLATensorPtr& grad_output,
   XLATensorPtr z = tensor_methods::exp(scaled_input);
   XLATensorPtr one_vec =
       tensor_methods::full_like(z, 1, z->GetDevice(), z->dtype());
+  XLATensorPtr div = GetValueOrThrow(
+      tensor_methods::div(z, tensor_methods::add(z, one_vec, 1)));
 
-  return tensor_methods::where(
-      tensor_methods::gt(scaled_input, threshold), grad_output,
-      tensor_methods::mul(
-          grad_output,
-          tensor_methods::div(z, tensor_methods::add(z, one_vec, 1))));
+  return tensor_methods::where(tensor_methods::gt(scaled_input, threshold),
+                               grad_output,
+                               tensor_methods::mul(grad_output, div));
 }
 
 XLATensorPtr Select(const XLATensorPtr& input, int64_t dim, int64_t index) {
@@ -223,8 +223,8 @@ XLATensorPtr EmbeddingDenseBackward(const XLATensorPtr& grad_output,
     XLATensorPtr grad_weights_scale =
         tensor_methods::index(counts, {indices_rank1}, 0);
     // Scale the value of the gradient by the histogram.
-    grad = tensor_methods::div(
-        grad, tensor_methods::unsqueeze(grad_weights_scale, 1));
+    grad = GetValueOrThrow(tensor_methods::div(
+        grad, tensor_methods::unsqueeze(grad_weights_scale, 1)));
   }
   // Don't accumulate gradients for indices which are equal with the given
   // padding_idx.