Improve error handling and error messages for uniform_.

Author: ysiraichi

test/test_operations.py

@@ -2367,6 +2367,19 @@ def test_isneginf_no_fallback(self):
     t = t.to(torch.float16)
     self._test_no_fallback(torch.isneginf, (t,))
 
+  def test_uniform__raises_error_on_invalid_range(self):
+    device = torch_xla.device()
+    a = torch.empty(5, 5, device=device)
+    from_ = 5.
+    to_ = 2.
+
+    try:
+      a.uniform_(from_, to_)
+    except RuntimeError as e:
+      expected_error = (
+          "uniform_(): expected `from` (5) to be smaller or equal `to` (2).")
+      self.assertEqual(str(e), expected_error)
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/aten_xla_type.cpp

@@ -3905,8 +3905,9 @@ at::Tensor& XLANativeFunctions::uniform_(
     return at::native::call_fallback_fn<&xla_fallback, ATEN_OP(uniform_)>::call(
         self, from, to, generator);
   }
-  XLA_ASSIGN_OR_THROW(XLATensorPtr xla_self, bridge::GetXlaTensor(self));
-  tensor_methods::uniform_(xla_self, from, to);
+  XLA_ASSIGN_OR_THROW(absl_nonnull XLATensorPtr xla_self,
+                      bridge::GetXlaTensor(self));
+  XLA_THROW_IF_ERROR(tensor_methods::uniform_(xla_self, from, to));
   return self;
 }
 

torch_xla/csrc/tensor_methods.cpp

@@ -14,6 +14,7 @@
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_split.h"
+#include "status.h"
 #include "torch_xla/csrc/LazyIr.h"
 #include "torch_xla/csrc/aten_xla_bridge.h"
 #include "torch_xla/csrc/data_ops.h"
@@ -506,6 +507,15 @@ absl::Status CheckMMMatrixSizesAreCompatible(const XLATensorPtr& mat1,
   return absl::OkStatus();
 }
 
+absl::Status CheckUniformRangeIsValid(double from, double to) {
+  if (from > to) {
+    return XLA_ERROR_WITH_LOCATION(absl::InvalidArgumentError(
+        absl::StrCat("uniform_(): expected `from` (", from,
+                     ") to be smaller or equal `to` (", to, ").")));
+  }
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 //////////////////////////////////////////////////////////////////////////////
@@ -3653,15 +3663,16 @@ std::vector<XLATensorPtr> unbind(const XLATensorPtr& input, int64_t dim) {
   return slices;
 }
 
-void uniform_(XLATensorPtr& input, double from, double to) {
-  XLA_CHECK_LE(from, to);
-  auto input_shape = input->shape();
+absl::Status uniform_(XLATensorPtr& input, double from, double to) {
+  XLA_RETURN_IF_ERROR(CheckUniformRangeIsValid(from, to));
+  xla::Shape input_shape = input->shape();
   input->SetInPlaceIrValue(torch_xla::MakeNode<Uniform>(
       XLAGraphExecutor::Get()->GetIrValueForScalar(
-          from, input_shape.get().element_type(), input->GetDevice()),
+          from, input_shape.element_type(), input->GetDevice()),
       XLAGraphExecutor::Get()->GetIrValueForScalar(
-          to, input_shape.get().element_type(), input->GetDevice()),
+          to, input_shape.element_type(), input->GetDevice()),
       XLAGraphExecutor::Get()->GetRngSeed(input->GetDevice()), input_shape));
+  return absl::OkStatus();
 }
 
 XLATensorPtr unsqueeze(const XLATensorPtr& input, int64_t dim) {

torch_xla/csrc/tensor_methods.h

@@ -986,7 +986,7 @@ std::tuple<XLATensorPtr, XLATensorPtr> triangular_solve(
 // removed.
 std::vector<XLATensorPtr> unbind(const XLATensorPtr& input, int64_t dim);
 
-void uniform_(XLATensorPtr& input, double from, double to);
+absl::Status uniform_(XLATensorPtr& input, double from, double to);
 
 // Insert a dimension of size one at the specified position.
 XLATensorPtr unsqueeze(const XLATensorPtr& input, int64_t dim);