Fix error messages.

Author: ysiraichi

test/stablehlo/test_stablehlo_custom_call.py

@@ -180,27 +180,6 @@ def get_graph(gm: torch.fx.GraphModule, _):
     self.assertNotIn("place_to_host", bw.code)
     self.assertNotIn("place_to_device", bw.code)
 
-  def test_stablehlo_custom_call_tracing_error(self):
-
-    def test_no_input():
-      stablehlo_custom_call(tuple(), "custom_op_target", [[1]], [torch.int8])
-
-    self.assertExpectedRaisesInline(
-        exc_type=RuntimeError,
-        callable=test_no_input,
-        expect="""custom_call(custom_op_target): expected at least 1 input tensor."""
-    )
-
-    def test_output_properties_mismatch():
-      input = torch.rand(10, device=torch_xla.device())
-      stablehlo_custom_call((input,), "custom_op_target", [[1], [1]], [torch.int8])
-
-    self.assertExpectedRaisesInline(
-        exc_type=RuntimeError,
-        callable=test_output_properties_mismatch,
-        expect="""custom_call(custom_op_target): expected the given output shapes (size=2) to be of the same size as the given output dtypes (size=1)."""
-    )
-
 
 if __name__ == "__main__":
 

test/test_ops_error_message.py

@@ -1,3 +1,4 @@
+from typing import Callable
 import expecttest
 import os
 import torch
@@ -357,6 +358,56 @@ def gen_test_fn(kernel_size=[2, 2, 2], stride=[], padding=[0]):
         expect="""avg_pool3d(): expected argument padding [1, 2] (size: 2) to have size of 3."""
     )
 
+  def _get_custom_call_properties(self, mode):
+    match mode:
+      case "tpu":
+        return (torch_xla._XLAC._xla_tpu_custom_call, "", [])
+      case "stablehlo":
+        return (torch_xla._XLAC._xla_custom_call, "custom_op_target",
+                [False, "", 0, {}])
+
+    self.fail(f"expected `mode` ({mode}) to be either of ['tpu', 'stablehlo'].")
+
+  def _gen_custom_call_no_input(self, mode):
+    lib_custom_call, payload, args = self._get_custom_call_properties(
+        mode)  # type: ignore[attr-defined]
+    return lambda: lib_custom_call([], payload, [[1]], [torch.int8], *args)
+
+  def _gen_custom_call_output_properties_size_mismatch(self, mode):
+    lib_custom_call, payload, args = self._get_custom_call_properties(
+        mode)  # type: ignore[attr-defined]
+    input = torch.rand(10, device=torch_xla.device())
+    return lambda: lib_custom_call(
+        (input,), payload, [[1], [1]], [torch.int8], *args)
+
+  def test_stablehlo_custom_call(self):
+
+    self.assertExpectedRaisesInline(
+        exc_type=RuntimeError,
+        callable=self._gen_custom_call_no_input("stablehlo"),
+        expect="""custom_call(custom_op_target): expected at least 1 input tensor."""
+    )
+
+    self.assertExpectedRaisesInline(
+        exc_type=RuntimeError,
+        callable=self._gen_custom_call_output_properties_size_mismatch(
+            "stablehlo"),
+        expect="""custom_call(custom_op_target): expected the given output shapes (size=2) to be of the same size as the given output dtypes (size=1)."""
+    )
+
+  def test_tpu_custom_call(self):
+
+    self.assertExpectedRaisesInline(
+        exc_type=RuntimeError,
+        callable=self._gen_custom_call_no_input("tpu"),
+        expect="""tpu_custom_call(): expected at least 1 input tensor.""")
+
+    self.assertExpectedRaisesInline(
+        exc_type=RuntimeError,
+        callable=self._gen_custom_call_output_properties_size_mismatch("tpu"),
+        expect="""tpu_custom_call(): expected the given output shapes (size=2) to be of the same size as the given output dtypes (size=1)."""
+    )
+
 
 if __name__ == "__main__":
   unittest.main()

torch_xla/csrc/tensor_methods.cpp

@@ -645,34 +645,51 @@ absl::Status CheckUniformRangeIsValid(double from, double to) {
   return absl::OkStatus();
 }
 
+// This check is used for both `custom_call()` and `tpu_custom_call()`.
+//
+// The `target` parameter is `std::nullopt` whenever it's being called from
+// a `tpu_custom_call()` context.
 absl::Status CheckCustomCallNonEmptyInputs(
     const std::vector<absl_nonnull XLATensorPtr>& inputs,
-    const std::string& target) {
+    const std::optional<std::string>& target) {
   if (inputs.empty()) {
-    return XLA_ERROR_WITH_LOCATION(absl::InvalidArgumentError(absl::StrCat(
-        "custom_call(", target, "): expected at least 1 input tensor.")));
+    std::string op = target.has_value()
+                         ? absl::StrCat("custom_call(", *target, ")")
+                         : "tpu_custom_call";
+    return XLA_ERROR_WITH_LOCATION(absl::InvalidArgumentError(
+        absl::StrCat(op, ": expected at least 1 input tensor.")));
   }
   return absl::OkStatus();
 }
 
+// This check is used for both `custom_call()` and `tpu_custom_call()`.
+//
+// The `target` parameter is `std::nullopt` whenever it's being called from
+// a `tpu_custom_call()` context.
 absl::Status CheckCustomCallOutputPropertiesSize(
     const std::vector<std::vector<int64_t>>& output_shapes,
     const std::vector<at::ScalarType>& output_dtypes,
-    const std::string& target) {
+    const std::optional<std::string>& target) {
   if (output_shapes.size() != output_dtypes.size()) {
+    std::string op = target.has_value()
+                         ? absl::StrCat("custom_call(", *target, ")")
+                         : "tpu_custom_call()";
     return XLA_ERROR_WITH_LOCATION(absl::InvalidArgumentError(absl::StrCat(
-        "custom_call(", target,
-        "): expected the given output shapes (size=", output_shapes.size(),
+        op, ": expected the given output shapes (size=", output_shapes.size(),
         ") to be of the same size as the given output dtypes (size=",
         output_dtypes.size(), ").")));
   }
   return absl::OkStatus();
 }
 
+// This check is used for both `custom_call()` and `tpu_custom_call()`.
+//
+// The `target` parameter is `std::nullopt` whenever it's being called from
+// a `tpu_custom_call()` context.
 template <class F>
 absl::StatusOr<std::vector<absl_nonnull XLATensorPtr>> CustomCallImpl(
     const std::vector<absl_nonnull XLATensorPtr>& inputs,
-    const std::string& target,
+    const std::optional<std::string>& target,
     const std::vector<std::vector<int64_t>>& output_shapes,
     const std::vector<at::ScalarType>& output_dtypes, F&& make_node) {
   XLA_RETURN_IF_ERROR(CheckCustomCallNonEmptyInputs(inputs, target));
@@ -995,15 +1012,15 @@ absl::StatusOr<std::vector<absl_nonnull XLATensorPtr>> tpu_custom_call(
     const std::vector<at::ScalarType>& output_dtypes) {
   XLA_ASSIGN_OR_RETURN(
       std::vector<absl_nonnull XLATensorPtr> outputs,
-      CustomCallImpl(inputs, payload, output_shapes, output_dtypes,
-                     /* make_node= */
-                     [&](const std::vector<torch::lazy::Value>& values,
-                         const std::vector<xla::Shape>& output_xla_shapes) {
-                       return torch_xla::MakeNode<TpuCustomCall>(
-                           values,
-                           xla::ShapeUtil::MakeTupleShape(output_xla_shapes),
-                           payload);
-                     }));
+      CustomCallImpl(
+          inputs, /* target= */ std::nullopt, output_shapes, output_dtypes,
+          /* make_node= */
+          [&](const std::vector<torch::lazy::Value>& values,
+              const std::vector<xla::Shape>& output_xla_shapes) {
+            return torch_xla::MakeNode<TpuCustomCall>(
+                values, xla::ShapeUtil::MakeTupleShape(output_xla_shapes),
+                payload);
+          }));
 
   XLAGraphExecutor* graph_executor = XLAGraphExecutor::Get();
   if (graph_executor->UseEagerMode()) {