Repository: pytorch/torchdynamo

Issue:
TorchDynamo is incompatible with einops.repeat
Hi,

We are doing some functionality validation and performance benchmark for TorchDynamo, and we found that TorchDynamo is incompatible with "einops" (pip package)'s "repeat" function (or maybe also some other functions in einops). For all models that call einops.repeat in nn.Module model definition (e.g. vision transformers models from https://github.com/lucidrains/vit-pytorch), the models can run with regular PyTorch but failed when using TorchDynamo with simply a FX backend, like the example shown in TorchDynamo homepage README.

Reproduction code (test.py):
```
import torch
from torch import nn
from einops import repeat
import argparse

parser = argparse.ArgumentParser(description='')
parser.add_argument('--regular', action='store_true')
parser.add_argument('--torchdynamo', action='store_true')
args = parser.parse_args()

class Net(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        repeat(x, 'h w -> h w c', c=3).shape # an einops.repeat() sentence
        return x

X = torch.randn(30, 40)
model = Net()

model.eval()

# Regular
if args.regular:
    model(X)

# Torchdynamo
if args.torchdynamo:
    import torchdynamo
    from typing import List
    
    def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
        return gm.forward
    
    with torchdynamo.optimize(my_compiler), torch.no_grad():
        model(X)
```
For this code snippet, ```python test.py --regular``` will pass but ```python test.py --torchdynamo``` will fail.

Error Log:
```
Traceback (most recent call last):
  File "/tmp/kyao/torchdynamo/test_einops.py", line 37, in <module>
    model(X)
  File "/mnt/local_disk/miniconda3/envs/ipexnew/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/tmp/kyao/torchdynamo/test_einops.py", line 15, in forward
    def forward(self, x):
  File "/mnt/local_disk/miniconda3/envs/ipexnew/lib/python3.9/site-packages/einops-0.4.1-py3.9.egg/einops/einops.py", line 490, in repeat
    def repeat(tensor: Tensor, pattern: str, **axes_lengths) -> Tensor:
  File "/mnt/local_disk/miniconda3/envs/ipexnew/lib/python3.9/site-packages/einops-0.4.1-py3.9.egg/einops/einops.py", line 410, in reduce
    return _apply_recipe(recipe, tensor, reduction_type=reduction)
  File "/mnt/local_disk/miniconda3/envs/ipexnew/lib/python3.9/site-packages/einops-0.4.1-py3.9.egg/einops/einops.py", line 229, in _apply_recipe
    def _apply_recipe(recipe: TransformRecipe, tensor: Tensor, reduction_type: Reduction) -> Tensor:
  File "/mnt/local_disk/miniconda3/envs/ipexnew/lib/python3.9/site-packages/einops-0.4.1-py3.9.egg/einops/_backends.py", line 22, in get_backend
    def get_backend(tensor) -> 'AbstractBackend':
  File "/mnt/local_disk/miniconda3/envs/ipexnew/lib/python3.9/site-packages/einops-0.4.1-py3.9.egg/einops/_backends.py", line 22, in get_backend
    def get_backend(tensor) -> 'AbstractBackend':
IndexError: list index out of range
```


Discussion:
Thanks for reporting this and providing a repro.  I'd expect just a graph break here (which falls back to eager), so clearly something is going wrong.

@yanboliang do you want to take a look at this one?
A smaller repro, that might make things easier. It seems to get the same error: 

```
import torch
from torch import nn
from einops import repeat, _backends
import argparse


parser = argparse.ArgumentParser(description='')
parser.add_argument('--regular', action='store_true')
parser.add_argument('--torchdynamo', action='store_true')
args = parser.parse_args()

X = torch.randn(30, 40)

# Regular
if args.regular:
    backend = _backends.get_backend(X)

# Torchdynamo
if args.torchdynamo:
    import torchdynamo
    from typing import List
    
    def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
        print("gm")
        gm.graph.print_tabular()
        return gm.forward
    
    with torchdynamo.optimize(my_compiler), torch.no_grad():
        backend = _backends.get_backend(X)
```

````
Traceback (most recent call last):
  File "test.py", line 29, in <module>
    backend = _backends.get_backend(X)
  File "/Users/voz/miniconda3/envs/torchdynamo/lib/python3.8/site-packages/einops/_backends.py", line 22, in get_backend
    def get_backend(tensor) -> 'AbstractBackend':
  File "/Users/voz/miniconda3/envs/torchdynamo/lib/python3.8/site-packages/einops/_backends.py", line 22, in get_backend
    def get_backend(tensor) -> 'AbstractBackend':
  File "/Users/voz/miniconda3/envs/torchdynamo/lib/python3.8/site-packages/einops/_backends.py", line 22, in get_backend
    def get_backend(tensor) -> 'AbstractBackend':
  [Previous line repeated 1 more time]
IndexError: list index out of range

@voznesenskym if you want to take a crack at it feel free to assign to yourself.  Just coordinate with @yanboliang, so there is no duplicated work.
Thanks @jansel. Will take a look and see what I can do :)

Resolving patch:
diff --git a/torchdynamo/variables/user_defined.py b/torchdynamo/variables/user_defined.py
--- a/torchdynamo/variables/user_defined.py
+++ b/torchdynamo/variables/user_defined.py
@@ -46,6 +46,32 @@ def var_getattr(self, tx, name: str) -> "VariableTracker":
 
         return super(UserDefinedClassVariable, self).var_getattr(tx, name)
 
+    def call_method(
+        self,
+        tx,
+        name,
+        args: "List[VariableTracker]",
+        kwargs: "Dict[str, VariableTracker]",
+    ) -> "VariableTracker":
+        if (
+            name == "__subclasses__"
+            and len(args) == 0
+            and not kwargs
+            and "__subclasses__" not in self.value.__dict__
+        ):
+            options = VariableTracker.propagate(self, args, kwargs.values())
+            options["mutable_local"] = MutableLocal()
+            subs_as_vars: List[VariableTracker] = list()
+            for sub in self.value.__subclasses__():
+                source = AttrSource(tx.import_source(sub.__module__), sub.__name__)
+                subs_as_vars.append(
+                    variables.UserDefinedClassVariable(sub, source=source)
+                )
+
+            return variables.ListVariable(subs_as_vars, **options)
+
+        return super().call_method(tx, args, kwargs)
+
     def call_function(
         self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
     ) -> "VariableTracker":


Test patch:
diff --git a/tests/test_misc.py b/tests/test_misc.py
--- a/tests/test_misc.py
+++ b/tests/test_misc.py
@@ -1223,6 +1223,83 @@ def test_const_dict_variable_python_type(self):
             collections.OrderedDict,
         )
 
+    def test_builtin_subclasses_as_method_on_class_type(self):
+        class Foo:
+            def __init__(name):
+                self.ame_ = name
+
+            def get_name(self):
+                return "Foo " + self.name_
+
+        class Bar(Foo):
+            def __init__(name):
+                self.name_ = name
+
+            def get_name(self):
+                return "Bar " + self.name_
+
+        class Baz(Foo):
+            def __init__(name):
+                self.name_ = name
+
+            def get_name(self):
+                return "Baz " + self.name_
+
+        subs_of_foo_reg = Foo.__subclasses__()
+
+        counter = CompileCounter()
+
+        @torchdynamo.optimize_assert(counter)
+        def fn():
+            return Foo.__subclasses__()
+
+        subs_of_foo_optim = fn()
+
+        self.assertEqual(len(subs_of_foo_reg), 2)
+        self.assertEqual(subs_of_foo_reg, subs_of_foo_optim)
+
+    def test_builtin_subclasses_as_method_on_var(self):
+        class Foo:
+            def __init__(name):
+                self.name_ = name
+
+            def get_name(self):
+                return "Foo " + self.name_
+
+        class Bar(Foo):
+            def __init__(name):
+                self.name_ = name
+
+            def get_name(self):
+                return "Bar " + self.name_
+
+        class Baz(Bar):
+            def __init__(name):
+                self.name_ = name
+
+            def get_name(self):
+                return "Baz " + self.name_
+
+        subs_of_foo_reg = Foo.__subclasses__()
+        sub_of_foo_subclass_var_reg = subs_of_foo_reg[0].__subclasses__()
+
+        sub_of_foo_subclass_var_optim = list()
+        counter = CompileCounter()
+
+        @torchdynamo.optimize_assert(counter)
+        def fn():
+            return Foo.__subclasses__()
+
+        @torchdynamo.optimize_assert(counter)
+        def fn_single(subs_of_foo_optim):
+            return subs_of_foo_optim[0].__subclasses__()
+
+        subs_of_foo_optim = fn()
+        sub_of_foo_subclass_var_optim = fn_single(subs_of_foo_optim)
+
+        self.assertEqual(len(sub_of_foo_subclass_var_optim), 1)
+        self.assertEqual(sub_of_foo_subclass_var_optim, sub_of_foo_subclass_var_reg)
+
     def test_enum_no_graphbreaks(self):
         class Foo(enum.Enum):
             FOO = 0