Repository: pytorch
Issue: [dynamo] calling __torch_function__ with dynamically created subclass of torch.Tensor fails compilation
Labels: triaged,module: __torch_function__,module: dynamo,dynamo-must-fix,dynamo-tensor-subclasses

### 🐛 Describe the bug

I apologize in advance to whoever will take care of this bug, as it must be the most obscure corner case of all times.
But I swear, this is breaking our production code right now.

So here we go, the following code snippet breaks:

```python
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = MyClass.__subclasses__()
        types_ = tuple(
            torch.Tensor if t in subclasses else t for t in [type(self)]
        )
        return torch.Tensor.__torch_function__(torch.abs, types_, torch.rand(1), {})


def create_subclass(parents):
    class MySubClass(*parents):
        ...

    return MySubClass


def call_foo(x):
    return x.foo()


cls = create_subclass((MyClass,))

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # fails
```
```terminal
Traceback (most recent call last):
  File "/home/johannes/.config/JetBrains/PyCharmCE2023.2/scratches/scratch_189.py", line 27, in <module>
    torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # fails
  File "/home/johannes/Documents/jina/docarrayv2/venv/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/home/johannes/.config/JetBrains/PyCharmCE2023.2/scratches/scratch_189.py", line 21, in call_foo
    return x.foo()  # works
  File "/home/johannes/.config/JetBrains/PyCharmCE2023.2/scratches/scratch_189.py", line 5, in foo
    def foo(self):
AttributeError: module '__main__' has no attribute 'MySubClass'
```

This does **not** happen if *any one* of the following criteria is met:
- the subclass of `MyClass` is not created dynamically
- `MyClass.__subclasses__()` is not used for `subclasses`
- `MyClass.__subclasses__()` _is_ used for `subclasses`,  but `subclasses` is not used
- `torch.Tensor.__torch_function__()` is not called

<details>
<summary>Code snippets for each of the above</summary>

All of the below produce **no error**:

**1. The subclass of `MyClass` is not created dynamically**
```python
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = MyClass.__subclasses__()
        types_ = tuple(
            torch.Tensor if t in subclasses else t for t in [type(self)]
        )
        return torch.Tensor.__torch_function__(torch.abs, types_, torch.rand(1), {})



def call_foo(x):
    return x.foo()

class MySubClass(MyClass):
    ...

cls = MySubClass

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # works
```

**2. `MyClass.__subclasses__()` is not used for `subclasses`**
```python
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = [MyClass]
        types_ = tuple(
            torch.Tensor if t in subclasses else t for t in [type(self)]
        )
        return torch.Tensor.__torch_function__(torch.abs, types_, torch.rand(1), {})


def create_subclass(parents):
    class MySubClass(*parents):
        ...

    return MySubClass


def call_foo(x):
    return x.foo()


cls = create_subclass((MyClass,))

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # works
```
**3. `MyClass.__subclasses__()` _is_ used for `subclasses`, but `subclasses` is not used**

```python
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = MyClass.__subclasses__()
        types_ = tuple(
            torch.Tensor for t in [type(self)]
        )
        return torch.Tensor.__torch_function__(torch.abs, types_, torch.rand(1), {})


def create_subclass(parents):
    class MySubClass(*parents):
        ...

    return MySubClass


def call_foo(x):
    return x.foo()


cls = create_subclass((MyClass,))

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # works
```

**4. `torch.Tensor.__torch_function__()` is not called**

```python
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = MyClass.__subclasses__()
        types_ = tuple(
            torch.Tensor if t in subclasses else t for t in [type(self)]
        )
        return types_


def create_subclass(parents):
    class MySubClass(*parents):
        ...

    return MySubClass


def call_foo(x):
    return x.foo()


cls = create_subclass((MyClass,))

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # works
```
------------------------
</details>

So from the above __one could conclude that the problem is passing `MyClass.__subclasses__()` to `torch.Tensor.__torch_function__()`, right?__

Well, here where it gets weird (at least to me). __The error happens even if `types_` is never passed to `torch.Tensor.__torch_function__()`, as long as it is assigned somewhere!__

```python
# FAILS THE SAME WAY AS THE CODE ABOVE
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = MyClass.__subclasses__()
        types_ = tuple(
            torch.Tensor if t in subclasses else t for t in [type(self)]
        )  # this is never used anywhere!
        types_ = (torch.Tensor,)  # because it is overwritten here
        return torch.Tensor.__torch_function__(torch.abs, types_, torch.rand(1), {})


def create_subclass(parents):
    class MySubClass(*parents):
        ...

    return MySubClass


def call_foo(x):
    return x.foo()


cls = create_subclass((MyClass,))

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # fails
```

Commenting out the **unused (!)** line fixes it:
```python
# THIS WORKS!
import torch


class MyClass(torch.Tensor):
    def foo(self):
        subclasses = MyClass.__subclasses__()
        # types_ = tuple(
        #     torch.Tensor if t in subclasses else t for t in [type(self)]
        # )  # commenting out the unused line
        types_ = (torch.Tensor,)
        return torch.Tensor.__torch_function__(torch.abs, types_, torch.rand(1), {})


def create_subclass(parents):
    class MySubClass(*parents):
        ...

    return MySubClass


def call_foo(x):
    return x.foo()


cls = create_subclass((MyClass,))

call_foo(cls(torch.rand(1000, 1000)))  # works
torch.compile(call_foo, backend='eager')(cls(torch.rand(1000, 1000)))  # fails
```
```terminal
Process finished with exit code 0
```

So either I am missing something very obvious, or I just brought you a monstrosity of a corner case bug.
In the latter case, sorry, and salute! 🫡

### Versions

Collecting environment information...
PyTorch version: 2.0.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 10.5.0-1ubuntu1~20.04) 10.5.0
Clang version: Could not collect
CMake version: version 3.27.0
Libc version: glibc-2.31

Python version: 3.8.10 (default, May 26 2023, 14:05:08)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-78-generic-x86_64-with-glibc2.29
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   39 bits physical, 48 bits virtual
CPU(s):                          8
On-line CPU(s) list:             0-7
Thread(s) per core:              2
Core(s) per socket:              4
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           140
Model name:                      11th Gen Intel(R) Core(TM) i7-1195G7 @ 2.90GHz
Stepping:                        2
CPU MHz:                         2469.934
CPU max MHz:                     5000,0000
CPU min MHz:                     400,0000
BogoMIPS:                        5836.80
Virtualization:                  VT-x
L1d cache:                       192 KiB
L1i cache:                       128 KiB
L2 cache:                        5 MiB
L3 cache:                        12 MiB
NUMA node0 CPU(s):               0-7
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l2 invpcid_single cdp_l2 ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves split_lock_detect dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid movdiri movdir64b fsrm avx512_vp2intersect md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] mypy==1.3.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.4
[pip3] torch==2.0.1+cpu
[conda] Could not collect

cc @hameerabbasi @rgommers @ezyang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @peterbell10 @aakhundov @anijain2305 @Xia-Weiwen @ipiszy