[custom_op] better error message on no returns (#129896)

zou3519 · pytorchmergebot · commit 872d972e4159 · 2024-07-02T23:34:23.000Z
I run into this a lot. I can imagine that it would look opaque to users, so made it more friendly Old error message: "ValueError: infer_schema(func): Return has unsupported type <class 'inspect._empty'>." Test Plan: - new tests Pull Request resolved: #129896 Approved by: https://github.com/yushangdi
diff --git a/test/test_custom_ops.py b/test/test_custom_ops.py
@@ -595,6 +595,15 @@ def blah8(x, *, y=1):
         def blah9(x, *, y):
             pass
 
+    def test_infer_schema_no_return(self):
+        with self.assertRaisesRegex(
+            ValueError, "No return type annotation was provided. Please add one."
+        ):
+
+            @torch.library.custom_op("mylib::foo", mutates_args={})
+            def foo(x: torch.Tensor, y: int):
+                return x * y
+
     def test_infer_schema_supported(self):
         def a(x: Tensor) -> Tensor:
             return torch.empty([])
diff --git a/torch/_library/infer_schema.py b/torch/_library/infer_schema.py
@@ -187,6 +187,9 @@ def parse_return(annotation, error_fn):
     if annotation is None:
         return "()"
 
+    if annotation is inspect.Parameter.empty:
+        error_fn("No return type annotation was provided. Please add one.")
+
     origin = typing.get_origin(annotation)
     if origin is not tuple:
         if annotation not in SUPPORTED_RETURN_TYPES.keys():
