[jit] support `isinstance(foo, dict)`

[suo]

Today doing something like:

@torch.jit.script
def foo(s: Any):
    if isinstance(s, dict):
        print(s.items())

gives:

RuntimeError: Unknown type name 'dict':

And capitalizing "dict" will produce a segfault. Yikes!

A related issue is that doing something like:

isinstance(foo, Dict[str, Tensor])

is not actually possible in Python, since the type annotations are mostly static constructs that the interpreter doesn't know anything about.

A potential compromise is to have isinstance(foo, dict) refine the type of foo to Dict[Any, Any]. And then the user can further type refine to get an actual things they want. So:

if isinstance(foo, dict):
    for k, v in foo.items():
        assert isinstance(k, str)
        assert isinstance(v, Tensor)

Another potential solution is to have a torch.jit.isinstance(foo, Dict[str, tensor]) that boils down to isinstance(foo, dict) in Python but performs the correct type refinement in TS.

cc @suo @gmagogsfm

[jansel]

Not clear how to solve this, need to come up with design and run it by others.

[edqwerty10]

Hi pytorch team, just a question, are there plans to also implement a similar check with list, ie. isinstance(foo, list).
Currently, the following code breaks in script mode:
Code Snippet:

import torch
from typing import Any

class TestModule(torch.nn.Module):
    def __init__(self):
        super(TestModule, self).__init__()
    
    def call(self, input1: str, input2: str) -> str:
        return input1

    def forward(self, input: Any) -> None:
        if isinstance(input, list):
            for el in input:
                print(el)

scripted_module = torch.jit.script(TestModule())
scripted_module(["1", "2"])

Output

'list' object is not iterable:
  File "<ipython-input-4-1d53cd51275f>", line 13
    def forward(self, input: Any) -> None:
        if isinstance(input, list):
            for el in input:
            ~~~~~~~~~~~~~~~~
                print(el)
                ~~~~~~~~ <--- HERE

[yf225]

This is important for avoiding code duplication in PyPer models. Consider the following use case:

import torch
from typing import Any, List

class TestModule(torch.nn.Module):
    def __init__(self):
        super(TestModule, self).__init__()
    
    def call(self, input1: str, input2: str) -> str:
        return input1

    def forward(self, input: Any) -> None:
        if torch.jit.is_scripting() and isinstance(input, List[str]):
            # This can be a big code block in actual use cases
            for el in input:
                print(el)
        elif not torch.jit.is_scripting() and isinstance(input, list):
            # Code duplication!
            for el in input:
                print(el)
        else:
            raise Exception

module = TestModule()
scripted_module = torch.jit.script(module)

module(["1", "2"])
scripted_module(["1", "2"])

One theoretically plausible workaround which doesn't work:

def forward(self, input: Any) -> None:
    if (torch.jit.is_scripting() and isinstance(input, List[str])) or (
        not torch.jit.is_scripting() and isinstance(input, list)
    ):  # Unfortunately this doesn't work in scripting...
        for el in input:
            print(el)

So to avoid code duplication, we need isinstance(foo, dict) and isinstance(foo, list) to work in TorchScript. :)

[gmagogsfm]

I am pretty sure I am missing something, but would it be enough to encapsulate the logic into functions to avoid duplication?

To be clear, even with all code in a function, this is still inconvenient and we should improve the situation, just wanted to get more clarity. Thanks.

[gmagogsfm]

How do we solve the type covariance problem if we treat isinstance(foo, Dict) as isinstance(foo, Dict[Any, Any])? Without changing covariance rules, this would never evaluate to True. Therefore, it is really not a viable solution.

Creating torch.jit.isinstance is suboptimal as it requires users to change their code, but it does sidestep this problem. Thoughts?

[gmagogsfm]

Pointers to implementing torch.jit.isinstance if we were to take that route.

1. Add torch.jit.isinstance in python layer. In eager mode, it checks whether given object is Dict, then check all of its keys/values for subtype.

2. For scripting, a SpecialFormValue should be generated to represent a call to torch.jit.isinstance. It should be created here:
   

   pytorch/torch/csrc/jit/python/python_sugared_value.cpp

   Line 918 in cf48872

   return SpecialFormValue::create(prim::annotate);

3. torch.jit.isinstance should ultimately translate to something like prim::isinstance(foo, DictType[TensorType, StrType], we already support this operator.

Essentially, we are enhancing Python's isinstance, which doesn't support subtypes, by providing this API so that users can write same code for both eager and scripting mode.

Optionally but ideally, we should also support other container types like dict, optional, tuple.