Repository: pytorch/pytorch

Issue:
torch.jit.script segfault
## 🐛 Bug

Segfault on master (built Oct 4) with the following code (which is perhaps not expected to be scriptable):
```python
from typing import Any
import torch

class A:
    def __init__(self, t):
        self.t = t
    @staticmethod
    def f(a: torch.Tensor):
        return A(a + 1)

class B(A):
    def __init__(self, t):
        self.t = t + 10
    @staticmethod
    def f(a: torch.Tensor):
        return A(a + 1)

x = A(torch.tensor([3]))

def fun(x: Any):
    if isinstance(x, A):
        return A.f(x.t)
    else:
        return B.f(x.t)

print(torch.__version__)
sc = torch.jit.script(fun)
```


cc @ezyang @gchanan @zou3519 @gmagogsfm


Discussion:
I can repro, taking a look at it now.
Took an initial look, error is coming from an `try_ann_to_type` call and deep in pybind. Assigning to @SplitInfinity, who knows more about scripting classes.
The issue here is that `B` extends `A`, and so in `try_ann_to_type`, TorchScript sees a `__torch_script_class__` attribute on `B` that was attached to `A`, thinks that the class has already been scripted, and tries to look up the JIT type by name in the compilation unit type registry.
is this fixed?
#45940 is waiting for review with all comments addressed and all tests passing. I forgot to re-request review last week after pushing new changes, oops.

Resolving patch:
diff --git a/torch/jit/_state.py b/torch/jit/_state.py
--- a/torch/jit/_state.py
+++ b/torch/jit/_state.py
@@ -61,7 +61,6 @@ def enable():
 _script_classes = {}
 
 def _add_script_class(cls, name):
-    cls.__torch_script_class__ = True
     global _script_classes
     _script_classes[name] = cls
 
diff --git a/torch/jit/annotations.py b/torch/jit/annotations.py
--- a/torch/jit/annotations.py
+++ b/torch/jit/annotations.py
@@ -6,6 +6,7 @@
 from .._jit_internal import List, Tuple, is_tuple, is_list, Dict, is_dict, Optional, \
     is_optional, _qualified_name, Any, Future, is_future, is_ignored_fn
 from .._jit_internal import BroadcastingList1, BroadcastingList2, BroadcastingList3  # type: ignore
+from ._state import _get_script_class
 
 from torch._C import TensorType, TupleType, FloatType, IntType, \
     ListType, StringType, DictType, BoolType, OptionalType, ClassType, InterfaceType, AnyType, NoneType, \
@@ -316,16 +317,18 @@ def try_ann_to_type(ann, loc):
     if ann is torch.dtype:
         return IntType.get()  # dtype not yet bound in as its own type
     if inspect.isclass(ann) and issubclass(ann, enum.Enum):
-        if not hasattr(ann, "__torch_script_class__"):
+        qualified_name = _qualified_name(ann)
+        if _get_script_class(qualified_name) is None:
             torch.jit._script._recursive_compile_class(ann, loc)
         return EnumType(_qualified_name(ann), get_enum_value_type(ann, loc), list(ann))
     if inspect.isclass(ann):
-        if hasattr(ann, "__torch_script_class__"):
-            return ClassType(_qualified_name(ann))
+        qualified_name = _qualified_name(ann)
+        if _get_script_class(qualified_name) is not None:
+            return ClassType(qualified_name)
         ignored_builtin_classes = (torch.nn.Module, tuple, list, Exception)
         if torch._jit_internal.can_compile_class(ann) and not issubclass(ann, ignored_builtin_classes):
             torch.jit._script._recursive_compile_class(ann, loc)
-            return ClassType(_qualified_name(ann))
+            return ClassType(qualified_name)
 
     # Maybe resolve a NamedTuple to a Tuple Type
     def fake_rcb(key):


Test patch:
diff --git a/test/jit/test_class_type.py b/test/jit/test_class_type.py
--- a/test/jit/test_class_type.py
+++ b/test/jit/test_class_type.py
@@ -6,6 +6,7 @@
 import torch
 import torch.nn as nn
 from torch.testing import FileCheck
+from typing import Any
 
 # Make the helper files in test/ importable
 pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
@@ -445,6 +446,39 @@ class Derived(Base):
                 def two(self, x):
                     return x + self.b + 2
 
+
+    def test_class_inheritance_implicit(self):
+        """
+        Test that inheritance is detected in
+        implicit scripting codepaths (e.g. try_ann_to_type).
+        """
+        class A:
+            def __init__(self, t):
+                self.t = t
+
+            @staticmethod
+            def f(a: torch.Tensor):
+                return A(a + 1)
+
+        class B(A):
+            def __init__(self, t):
+                self.t = t + 10
+
+            @staticmethod
+            def f(a: torch.Tensor):
+                return A(a + 1)
+
+        x = A(torch.tensor([3]))
+
+        def fun(x: Any):
+            if isinstance(x, A):
+                return A.f(x.t)
+            else:
+                return B.f(x.t)
+
+        with self.assertRaisesRegex(RuntimeError, "Tried to access nonexistent attribute or method"):
+            sc = torch.jit.script(fun)
+
     @unittest.skipIf(IS_SANDCASTLE, "Importing like this doesn't work in fbcode")
     def test_imported_classes(self):
         import jit._imported_class_test.foo
diff --git a/torch/testing/_internal/jit_utils.py b/torch/testing/_internal/jit_utils.py
--- a/torch/testing/_internal/jit_utils.py
+++ b/torch/testing/_internal/jit_utils.py
@@ -55,6 +55,7 @@ def do_input_map(fn, input):
 def clear_class_registry():
     torch._C._jit_clear_class_registry()
     torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
+    torch.jit._state._script_classes.clear()
 
 def get_execution_plan(graph_executor_state):
     execution_plans = list(graph_executor_state.execution_plans.values())