Repository: pytorch
Issue: TorchScript cannot handle declared module attributes of Python types
Labels: oncall: jit

## 🐛 Bug

If a module declares an annotated attribute with a type that is not TorchScript compatible, it cannot be `torch.jit.script`ed even if that attribute is never used in TorchScript.

## To Reproduce

```
In [1]: import torch

In [3]: class Thing:
   ...:     def __init__(self, a: int):
   ...:         self.a = a
   ...: 
   ...: 

In [5]: class M(torch.nn.Module):
   ...:     thing: Thing
   ...:     a: int
   ...:     def __init__(self, thing: Thing):
   ...:         super().__init__()
   ...:         self.thing = thing
   ...:         self.a = thing.a
   ...:     def forward(self, x):
   ...:         return self.a * x
   ...: 

In [7]: t = Thing(4)

In [8]: x = torch.ones(3)

In [9]: m = M(t)

In [10]: m(x)
Out[10]: tensor([4., 4., 4.])

In [11]: m_script = torch.jit.script(m)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-11-6db52b40ddf3> in <module>
----> 1 m_script = torch.jit.script(m)

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_script.py in script(obj, optimize, _frames_up, _rcb)
    940     if isinstance(obj, torch.nn.Module):
    941         obj = call_prepare_scriptable_func(obj)
--> 942         return torch.jit._recursive.create_script_module(
    943             obj, torch.jit._recursive.infer_methods_to_compile
    944         )

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_recursive.py in create_script_module(nn_module, stubs_fn, share_types)                                                                                                                    
    388     assert not isinstance(nn_module, torch.jit.RecursiveScriptModule)
    389     check_module_initialized(nn_module)
--> 390     concrete_type = get_module_concrete_type(nn_module, share_types)
    391     return create_script_module_impl(nn_module, concrete_type, stubs_fn)
    392 

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_recursive.py in get_module_concrete_type(nn_module, share_types)                                                                                                                          
    364     if share_types:
    365         # Look into the store of cached JIT types
--> 366         concrete_type = concrete_type_store.get_or_create_concrete_type(nn_module)
    367     else:
    368         # Get a concrete type directly, without trying to re-use an existing JIT

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_recursive.py in get_or_create_concrete_type(self, nn_module)
    305         types are re-used if possible.
    306         """
--> 307         concrete_type_builder = infer_concrete_type_builder(nn_module)
    308 
    309         nn_module_type = type(nn_module)

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_recursive.py in infer_concrete_type_builder(nn_module, share_types)                                                                                                                       
    269 
    270         # If we got here, this is a regular "data" attribute, Add it to the concrete type
--> 271         attr_type, inferred = infer_type(name, value)
    272         if attr_type.success():
    273             concrete_type_builder.add_attribute(name, attr_type.type(), False, False)

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_recursive.py in infer_type(name, item)
    121         inferred = False
    122         if name in class_annotations and class_annotations[name] != torch.nn.Module.__annotations__["forward"]:
--> 123             ann_to_type = torch.jit.annotations.ann_to_type(class_annotations[name], _jit_internal.fake_range())
    124             attr_type = torch._C.InferredType(ann_to_type)
    125         elif isinstance(item, torch.jit.Attribute):

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/annotations.py in ann_to_type(ann, loc)
    347 
    348 def ann_to_type(ann, loc):
--> 349     the_type = try_ann_to_type(ann, loc)
    350     if the_type is not None:
    351         return the_type

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/annotations.py in try_ann_to_type(ann, loc)
    337         ignored_builtin_classes = (torch.nn.Module, tuple, list, Exception)
    338         if torch._jit_internal.can_compile_class(ann) and not issubclass(ann, ignored_builtin_classes):
--> 339             torch.jit._script._recursive_compile_class(ann, loc)
    340             return ClassType(qualified_name)
    341 

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_script.py in _recursive_compile_class(obj, loc)
   1104     error_stack = torch._C.CallStack(_qual_name, loc)
   1105     rcb = _jit_internal.createResolutionCallbackForClassMethods(obj)
-> 1106     _compile_and_register_class(obj, rcb, _qual_name)
   1107 
   1108 CompilationUnit = torch._C.CompilationUnit

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/_script.py in _compile_and_register_class(obj, rcb, qualified_name)
     63 
     64 def _compile_and_register_class(obj, rcb, qualified_name):
---> 65     ast = get_jit_class_def(obj, obj.__name__)
     66     defaults = torch.jit.frontend.get_default_args_for_class(obj)
     67     torch._C._jit_script_class_compile(qualified_name, ast, defaults, rcb)

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/jit/frontend.py in get_jit_class_def(cls, self_name)
    178     properties = get_class_properties(cls, self_name)
    179 
--> 180     sourcelines, file_lineno, filename = get_source_lines_and_file(cls, torch._C.ErrorReport.call_stack())
    181     source = ''.join(sourcelines)
    182     dedent_src = dedent(source)

~/anaconda3/envs/sciml/lib/python3.8/site-packages/torch/_utils_internal.py in get_source_lines_and_file(obj, error_msg)
     52     filename = None  # in case getsourcefile throws
     53     try:
---> 54         filename = inspect.getsourcefile(obj)
     55         sourcelines, file_lineno = inspect.getsourcelines(obj)
     56     except OSError as e:

~/anaconda3/envs/sciml/lib/python3.8/inspect.py in getsourcefile(object)
    694     Return None if no way can be identified to get the source.
    695     """
--> 696     filename = getfile(object)
    697     all_bytecode_suffixes = importlib.machinery.DEBUG_BYTECODE_SUFFIXES[:]
    698     all_bytecode_suffixes += importlib.machinery.OPTIMIZED_BYTECODE_SUFFIXES[:]

~/anaconda3/envs/sciml/lib/python3.8/inspect.py in getfile(object)
    663             if getattr(module, '__file__', None):
    664                 return module.__file__
--> 665         raise TypeError('{!r} is a built-in class'.format(object))
    666     if ismethod(object):
    667         object = object.__func__

TypeError: <class '__main__.Thing'> is a built-in class
```

Removing the type declaration works fine:
```
In [12]: class M(torch.nn.Module):
    ...:     a: int
    ...:     def __init__(self, thing: Thing):
    ...:         super().__init__()
    ...:         self.thing = thing
    ...:         self.a = thing.a
    ...:     def forward(self, x):
    ...:         return self.a * x
    ...: 

In [13]: m = M(t)

In [14]: m_script = torch.jit.script(m)
```

## Expected behavior

Since `thing` is never used by a TorchScript method, it shouldn't matter whether it is decorated with a TorchScript compatible type. This will increase compatibility with modules that declare or store typed attributes for purposes besides TorchScript model inference.

## Environment

 - PyTorch Version (e.g., 1.0): 1.8.0
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source): n/a
 - Python version: 3.8.5
 - CUDA/cuDNN version: n/a
 - GPU models and configuration: n/a
 - Any other relevant information: n/a



cc @gmagogsfm