Repository: pytorch
Issue: [TorchScript] Rough edges discovered during type system rewrite
Labels: oncall: jit

Use this issue to document all the rough edges discovered during the (re-)writing of TorchScript type system (#50434 ).

- Behaviors of user-defined subclasses of torch.tensor are silently ignored. Instead it should report error when users try to subclass builtin nominal types  (opened a separate issue #51135 ) 

Test case: [test_subclass.py](https://github.com/gmagogsfm/TorchScriptStudy/blob/main/tensor_subclass.py). 

This test case overwrites `__len__` of `torch.Tensor` but the overwritten method is silently ignored
```
> python test-tensor-subclass.py
Eager and jit behavior mismatching: 123456 vs 2
```

Test case: [test-tensor-subclass2.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-tensor-subclass2.py)

This test case adds a new method `myNumber()` to a subclass of `torch.Tensor`, but when invoking the method, the method was not recognized (because subclasses were not compiled)
```
python test-tensor-subclass2.py 
Eager:  123456
Traceback (most recent call last):
  File "test-tensor-subclass2.py", line 18, in <module>
    scripted_fn = torch.jit.script(fn)
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 939, in script
    fn = torch._C._jit_script_compile(
RuntimeError: 
Tried to access nonexistent attribute or method 'myNumber' of type 'Tensor'.:
  File "test-tensor-subclass2.py", line 13
def fn(x: TensorSubClass):
    return x.myNumber()
           ~~~~~~~~~~ <--- HERE
``` 


- When module type constructor is invoked inside TorchScript, the error message is mystic. (opened a separate issue #51136)

Test case: [test-module.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-module.py)

```
python test-module.py 
Eager:  5
Traceback (most recent call last):
  File "test-module.bad.py", line 23, in <module>
    m = torch.jit.script(MyModel(2))
  File "/data/users/pengwu/pytorch/torch/jit/_script.py", line 901, in script
    qualified_name = _qualified_name(obj)
  File "/data/users/pengwu/pytorch/torch/_jit_internal.py", line 870, in _qualified_name
    raise RuntimeError("Could not get name of python class object")
RuntimeError: Could not get name of python class object
```

- interface types  do not allow to define constructors, but the error message complains about return type of `__init__` not annotated (opened separate issue #51137)

Test case: [test-interface.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-interface.py)
```
> python test-interface.py 
Traceback (most recent call last):
  File "test-interface.py", line 4, in <module>
    class TestInterface(torch.nn.Module):
  File "/data/users/pengwu/pytorch/torch/jit/_script.py", line 1043, in interface
    torch._C._jit_script_interface_compile(
RuntimeError: 
interface declarations must have a return type annotated.:
  File "test-interface.py", line 7
    def __init__(self):
    ~~~~~~~~~~~~~~~~~~~
        pass
        ~~~~ <--- HERE
```

- tensor printout format differs from eager mode (opened separate issue in #51138 )
Test case:  [test-tensor-printout-differ.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-tensor-printout-differ.py)
```
> python test-tensor-printout-differ.py 
tensor([1., 1., 1., 1., 1., 1.])
Eager: True
 1
 1
 1
 1
 1
 1
[ CPUFloatType{6} ]
TorchScript: True
```

- Early type-checking may yield more intuitive user errors (opened separate issue in #51139)

Test case: [test-class2.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-class2.py)

In this test case, a custom class defines a class attribute, which is not allowed. Instead of informing users that class attributes are not allowed, the error message is like this (treating the class attribute as an instance attribute):
```
RuntimeError: 
Tried to access nonexistent attribute or method 'name' of type '__torch__.MyClass'. Did you forget to initialize an attribute in __init__()?:
  File "test-class2.py", line 10
def fn(a: MyClass):
    return a.name
           ~~~~~~ <--- HERE
```

- Problems with `Any` type

Test case: [test-local-annotation2.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-local-annotation2.py)

In this test case, the type checker cannot infer a conditional expression into a type of `Any`, resulting in
```
python test-local-annotation2.py 
Traceback (most recent call last):
  File "test-local-annotation2.py", line 7, in <module>
    m = torch.jit.script(f)
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 939, in script
    fn = torch._C._jit_script_compile(
RuntimeError: 
if-expression's true branch has type int but false branch has type float:
  File "test-local-annotation2.py", line 4
def f(a: int, b: float, flag: bool):
    value: Any = a if flag else b
                 ~~~~~~~~~~~~~~~~ <--- HERE
    return value
```

Another test case: [test-local-annotation3.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-local-annotation3.py)

I tried to use if-then-else statement instead of conditional expressions, and found out that the value produced was wrong:
```
python test-local-annotation3.py 
Eager: 2
TorchScript: 1.0
```
- Non-intuitive error messages for user errors (opened separate issue in #51140)

Test case: [test-enum-mystic-error-message.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-enum-mystic-error-message.py) 

The test failed because `em_fn` was annotated as `@torch.script.jit`, but also supplied as input to `torch.script.jit(em_fn)`. The error message is not helpful to users. Instead, JIT should just tolerate `@torch.script.jit` instead of issuing errors. Similarly for modules that are annotated as `@torch.script.jit`but is also supplied as input to `torch.jit.trace(...)`
```
> python test-enum-mystic-error-message.py 
Traceback (most recent call last):
  File "test-enum-mystic-error-message.py", line 18, in <module>
    m = torch.jit.script(enum_fn)
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 936, in script
    ast = get_jit_def(obj, obj.__name__)
AttributeError: 'torch._C.ScriptFunction' object has no attribute '__name__'
```

Test case: [test-enum-annotated.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-enum-annotated.py)

The test failed because `enum` class is annotated as `@torch.jit.script` and is treated as a custom class
```
> python test-enum-annotated.py Traceback (most recent call last):
  File "test-enum-annotated.py", line 7, in <module>
    class Color(Enum):
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 913, in script
    raise RuntimeError(
RuntimeError: TorchScript classes must be new-style classes. Please inherit from 'object'.
```

Test case: [test-tuple.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-tuple.py)

Not exactly sure what is wrong with the example, perhaps it is related to decorators or passing a non-module object to `trace.jit.script()`?

```
> python test-tuple.py 
Traceback (most recent call last):
  File "test-tuple.py", line 14, in <module>
    m = torch.jit.script(MyClass(2, (1,2)))
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 901, in script
    qualified_name = _qualified_name(obj)
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/_jit_internal.py", line 793, in _qualified_name
    raise RuntimeError("Could not get name of python class object")
RuntimeError: Could not get name of python class object
```

Test case: [test-instanceclass.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-instanceclass.py)

This example passes an object of an arbitrary class to `trace.jit.script()`, but `trace.jit.script()` only accepts function objects or module objects. But the error message is not intuitive:
```
Traceback (most recent call last):
  File "test-instanceclass.py", line 14, in <module>
    m = torch.jit.script(MyClass(2, (1,2)))
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 901, in script
    qualified_name = _qualified_name(obj)
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/_jit_internal.py", line 793, in _qualified_name
    raise RuntimeError("Could not get name of python class object")
RuntimeError: Could not get name of python class object
```

A preferred error message should be something like this:
```
torch.jit.trace() expects function or module objects
```
- Type-checker too strict for Optional type initialization pattern

Test case: [test-optional.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-optional.py)
```
python test-optional.py 
Eager:  1
Traceback (most recent call last):
  File "test-optional.py", line 10, in <module>
    scripted_fn = torch.jit.script(fn)
  File "/home/pengwu/local/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/jit/_script.py", line 939, in script
    fn = torch._C._jit_script_compile(
RuntimeError: 
Variable 'b' previously has type None but is now being assigned to a value of type int
:
  File "test-optional.py", line 6
    b = None
    if a>100:
        b = 1
        ~ <--- HERE
    return b
```

- More intuitive error message for annotated `self` type

Test case: [test-self-annotation.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-self-annotation.py)

When annotating type for `self`, got the following error message:
```
python test-self-annotation.py 
Traceback (most recent call last):
  File "test-self-annotation.py", line 5, in <module>
    class MyClass(object):
  File "test-self-annotation.py", line 6, in MyClass
    def __init__(self: MyClass, x: int):
NameError: name 'MyClass' is not defined
```

A better error message is to tell users that `self` should not be annotated.

- Should report error when instance attributes are annotated outside `__init__`

Test case: [test-attr-annotation2.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-attr-annotation2.py)

Instance attributes should be annotated at class level, not inside `__init__` or any other methods. In this example, we annotate an instance attribute in an instance method and the compiler does not complain about it. This does not make sense because type schema of instance attributes are inferred when analyzing `__init__`. We should just ask programmers to remove such annotation to avoid confusion from reading the code.

- `Optional` type does not support `!= None` 

Test case: [test-optional-none.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-optional-none.py)

TS issues the following errors:
```
The original call is:
  File "test-optional-none.py", line 5
def fn(a: Optional[int], val: int):
    if a != None:
       ~~~~~~~~~ <--- HERE
        return a+val
    else:
```

It seems to support `a == None` or `a is not None` (when `a` is of Optional type), but not `a != None`. This asymmetry is unintuitive to users.

- Instance attribute annotation does not seem to take effect

Test case: [test-instance-attr-annotation.py](https://github.com/penguinwu/hodgepodge/blob/master/TorchScriptStudy/bad-examples/test-instance-attr-annotation.py)

In this example, instance attribute `x` is annotated to be `Optional[int]` at the class level, but the compiler does not seem to know that fact and produces the following error:
```
Wrong type for attribute assignment. Expected None but got int:
  File "test-instance-attr-annotation.py", line 11
        self.x = None
        if x > 0:
            self.x = x
            ~~~~~~~~~~ <--- HERE
```

- How to figure out which operators are supported for built-in types?

TorchScript only implements a subset of behaviors of the Python equivalent of its built-in types (e.g., primitive types, structural types, built-in classes). I have the feeling that primitive type support is probably on par w/ Python. But there may be some gap in terms of operators supported for structural types (e.g., list, dictionary) and built-in classes (e.g., tensors, ModuleList).

Such parity w/ Python is not only cumbersome to document, but it may also widen over time (e.g., torch library adding more operators to the built-in torch classes).

One thought is to allow users to print all the methods supported by a built-in type (e.g., `print(ModuleList.method_table)`).  This would be a tool to allow users to quickly figure out which methods are supported by a TorchScript type and any additional constraints. I believe we must have stored such meta-information internally for method resolution. 

cc @gmagogsfm