[JIT] Check mergeability of return types for functions that return Any

Meghan Lele · Meghan Lele · commit f368f2610f3b · 2020-07-29T17:37:54.000-07:00
**Summary** This commit modifies the type checking performed during IR generation so that it checks that all return types for a function marked as returning `Any` can be unified. Issue #41962 reported that the use of an `Any` return type in combination with different code paths returning values of different types causes a segmentation fault. During IR generation, the type of each return is checked against the annotated return type of the function, or unified with all other return types if there is no annotation. Because every type is a subtype of `Any`, this commit modifies the IR generation to do the latter (i.e. check if all return types can be unified) if the annotated return type is `Any`. **Test Plan** This commit adds a unit test that checks that an exception with an appropriate error message is thrown when a function and module with an annotated return type of `Any` is compiled and the possible return types cannot be unified. **Fixes** This commit fixes #41962. ghstack-source-id: bc3ca78 Pull Request resolved: #42259
diff --git a/test/test_jit_py3.py b/test/test_jit_py3.py
@@ -11,6 +11,7 @@
 import torch.testing._internal.jit_utils
 import torch.nn as nn
 import types
+import inspect
 
 class TestScriptPy3(JitTestCase):
     def test_joined_str(self):
@@ -533,6 +534,38 @@ def forward(self, x):
         # Check that ignored method is still intact.
         self.assertEqual(inp, n.ignored_method(inp))
 
+    def test_do_not_use_any_as_return_supertype(self):
+        """
+        Check that if a function is annotated as returning any,
+        all return paths are checked against each other (to see if
+        they have a common supertype) instead of against the
+        return type.
+        """
+        class Mod(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, inp: torch.Tensor) -> Any:
+                if inp.shape[0] == 1:
+                    return inp * inp
+                else:
+                    return 3
+
+        def any_function(inp: torch.Tensor) -> Any:
+            if inp.shape[0] == 1:
+                return inp * inp
+            else:
+                return 3
+
+        with self.assertRaisesRegex(RuntimeError, "Previous return statement returned a value of type Tensor but this return statement returns a value of type int"):
+            torch.jit.script(Mod())
+
+        with self.assertRaisesRegex(RuntimeError, "Previous return statement returned a value of type Tensor but this return statement returns a value of type int"):
+            torch.jit.script(any_function)
+
+        with self.assertRaisesRegex(RuntimeError, "Previous return statement returned a value of type Tensor but this return statement returns a value of type int"):
+            cu = torch.jit.CompilationUnit(inspect.getsource(any_function))
+
     def test_export_opnames_interface(self):
         global OneTwoModule
 
diff --git a/torch/csrc/jit/frontend/ir_emitter.cpp b/torch/csrc/jit/frontend/ir_emitter.cpp
@@ -956,7 +956,7 @@ struct to_ir {
     Value* result = emitExpr(stmt.expr());
     TypePtr result_type = def_stack_.back().declared_return_type_;
     // result type is annotated, every return must convert to that type
-    if (result_type) {
+    if (result_type && result_type != AnyType::get()) {
       // this guard skips implicit conversion from None -> Tensor for the return
       // type. otherwise forgetting a return a function returning a tensor will
       // cause a None to be converted to a tensor.
