Repository: pytorch/pytorch

Issue:
TorchDynamo does not support modules with custom __getitem__ implementation
### 🐛 Describe the bug

When writing a custom module which implements the __getitem__ method dynamo fails with AssertionError.

For more context on this (e.g. why a custom module for this) see: https://github.com/pytorch/pytorch/issues/71203
and
https://github.com/ludwig-ai/ludwig/blob/master/ludwig/features/feature_utils.py#L150

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire @jansel @agunapal 

### Error logs

```
Traceback (most recent call last):
  File "/Users/mreso/Projects/playground/ludwig/moduledict.py", line 31, in <module>
    cfoo(1)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 337, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 404, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 104, in _fn
    return fn(*args, **kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 262, in _convert_frame_assert
    return _compile(
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 324, in _compile
    out_code = transform_code_object(code, transform)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 445, in transform_code_object
    transformations(instructions, code_options)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 311, in transform
    tracer.run()
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1726, in run
    super().run()
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 576, in run
    and self.step()
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 540, in step
    getattr(self, inst.opname)(inst)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 342, in wrapper
    return inner_fn(self, inst)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 148, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 566, in call_function
    result = handler(tx, *args, **kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 790, in call_getitem
    return args[0].call_method(tx, "__getitem__", args[1:], kwargs)
  File "/Users/mreso/miniconda3/envs/pt2/lib/python3.10/site-packages/torch/_dynamo/variables/nn_module.py", line 402, in call_method
    assert type(module).__getitem__ in (
AssertionError: FeatureDict
from user code:
   File "/Users/mreso/Projects/playground/ludwig/moduledict.py", line 25, in forward
    return self.ld["key1"](x)
Set torch._dynamo.config.verbose=True for more information
You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True
```

### Minified repro

```python
import torch
class FeatureDict(torch.nn.Module):
    #Torch ModuleDict wrapper that permits keys with any name.
    def __init__(self):
        super().__init__()
        self.module_dict = torch.nn.ModuleDict()
    def __getitem__(self, key) -> torch.nn.Module:
        return self.module_dict[key]
    def __setitem__(self, key: str, module: torch.nn.Module) -> None:
        self.module_dict[key] = module

class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.ld = FeatureDict()
        
        self.ld["key1"] = torch.nn.Linear(1,1)
        
    def forward(self, x):
        return self.ld["key1"](x)

foo = Foo()
cfoo = torch.compile(foo)
cfoo(1)
```

### Versions

```
Collecting environment information...
PyTorch version: 2.0.0+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.21.3
Libc version: glibc-2.27

Python version: 3.9.16 (main, Mar  8 2023, 14:00:05)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-1097-aws-x86_64-with-glibc2.27
Is CUDA available: False
CUDA runtime version: 10.0.130
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
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
CPU(s):              16
On-line CPU(s) list: 0-15
Thread(s) per core:  2
Core(s) per socket:  8
Socket(s):           1
NUMA node(s):        1
Vendor ID:           GenuineIntel
CPU family:          6
Model:               85
Model name:          Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
Stepping:            7
CPU MHz:             3600.005
BogoMIPS:            5999.99
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            1024K
L3 cache:            36608K
NUMA node0 CPU(s):   0-15
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==2.0.0+cpu
[pip3] torchaudio==2.0.1+cpu
[pip3] torchvision==0.15.1+cpu
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] torch                     2.0.0+cpu                pypi_0    pypi
[conda] torchaudio                2.0.1+cpu                pypi_0    pypi
[conda] torchvision               0.15.1+cpu               pypi_0    pypi
```


Discussion:
We can support them by inlining the ```__getitem__``` function, but need to update guards, seems non-trivial. For a quick fix, we should remove the assertion and fallback to python for unsupported cases.

Resolving patch:
diff --git a/torch/_dynamo/variables/nn_module.py b/torch/_dynamo/variables/nn_module.py
--- a/torch/_dynamo/variables/nn_module.py
+++ b/torch/_dynamo/variables/nn_module.py
@@ -399,12 +399,28 @@ def gen_source(source, name):
             )
         elif name == "__getitem__":
             assert not kwargs and len(args) == 1
-            assert type(module).__getitem__ in (
+            builtin_supported = (
                 torch.nn.ModuleDict.__getitem__,
                 torch.nn.ModuleList.__getitem__,
                 torch.nn.ParameterList.__getitem__,
                 torch.nn.Sequential.__getitem__,
-            ), typestr(module)
+            )
+
+            if type(module).__getitem__ not in builtin_supported:
+                assert isinstance(args[0], variables.ConstantVariable), typestr(args[0])
+                key = args[0].as_python_constant()
+                assert isinstance(key, (str, int))
+                fn = getattr(module, name).__func__
+
+                assert isinstance(fn, types.FunctionType)
+
+                src = AttrSource(AttrSource(self.source, name), "__func__")
+                return tx.inline_user_function_return(
+                    variables.UserFunctionVariable(fn, source=src, **options),
+                    [self] + list(args),
+                    kwargs,
+                )
+
             assert self.source
 
             if isinstance(args[0], SliceVariable):


Test patch:
diff --git a/test/dynamo/test_modules.py b/test/dynamo/test_modules.py
--- a/test/dynamo/test_modules.py
+++ b/test/dynamo/test_modules.py
@@ -295,6 +295,31 @@ def forward(self, x):
         return x
 
 
+class CustomGetItemModuleList(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = torch.nn.ModuleList(
+            [
+                torch.nn.Linear(10, 10),
+                torch.nn.ReLU(),
+                torch.nn.Linear(10, 10),
+                torch.nn.ReLU(),
+            ]
+        )
+
+    def __getitem__(self, idx: int):
+        return self.layers[idx]
+
+    def __len__(self) -> int:
+        return len(self.layers)
+
+    def forward(self, x):
+        for i in range(len(self)):
+            x = self[i](x)
+
+        return x
+
+
 class ModuleDict(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -310,6 +335,23 @@ def forward(self, x):
         return x
 
 
+class CustomGetItemModuleDict(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = torch.nn.ModuleDict(
+            {
+                "0": torch.nn.Linear(10, 10),
+            }
+        )
+
+    def __getitem__(self, key: str) -> torch.nn.Module:
+        return self.layers[key]
+
+    def forward(self, x):
+        x = self["0"](x)
+        return x
+
+
 class TensorList(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -728,7 +770,9 @@ class NNModuleTests(torch._dynamo.test_case.TestCase):
     test_cfgmod = make_test(CfgModule())
     test_stringmember = make_test(StringMember())
     test_modulelist = make_test(ModuleList())
+    test_modulelist = make_test(CustomGetItemModuleList())
     test_moduledict = make_test(ModuleDict())
+    test_moduledict = make_test(CustomGetItemModuleDict())
     test_super1 = make_test(SuperModule())
     test_super2 = make_test(SuperModule2())
     test_super_class_method = make_test(SuperChildCallsClassMethod())