[JIT] Finish Resolving Python Types By Value

[eellison]

🐛 Bug

To recursively compile python programs, we need to recursively look up the python types that are used. We do this in many locations currently, but not all.

Something like:

class X(object):
    def __init__(self):
        pass

@torch.jit.script
def foo():
    return X()
    pass

Will correctly resolve X() and compile , but if you modify the function to be:

@torch.jit.script
def foo(out: X):
    return out

it will fail with Unknown type name 'X':. This is because in many locations in the compiler we try to resolve the code textually (e.g. here). The textual resolution is needed for compilation in C++, but when the python environment is available we should search it using resolveType here. We will need to handle cases that use . and [], such as List[int], and namespace.X.

This has been an issue for a number of users, and is a particular issue with torch.jit.Future. torch.jit.Future is the runtime type value, but if users use torch.jit.Future as an inline type annotation it will fail.

The first part of this issue was https://github.com/pytorch/pytorch/pull/29623/files. This issue may require multiple PRs to completely fix, and you may run into other issues when doing it.

cc @suo