[dynamo] AssertionError for custom iterable nn.Module
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

Using a custom iterable nn.Module causes an AssertionError when dynamo is enabled.
The use case here is creating an equivalent of nn.ParameterList, but for buffers, like in [this issue](https://github.com/pytorch/pytorch/issues/37386) or [this question](https://discuss.pytorch.org/t/why-no-nn-bufferlist-like-function-for-registered-buffer-tensor/18884/3)
The code hits an assertion in __dynamo/variables/nn_module.py_ line 99 (in _unpack_var_sequence_), which seems to explicitly disallow custom iterable modules.
Verified in nighly 2023.06.17


### Error logs

```
AssertionError                            Traceback (most recent call last)
[<ipython-input-2-79323aca98f5>](https://localhost:8080/#) in <cell line: 37>()
     35 x = Containing(BufferList((torch.rand((5,5)), torch.rand((5,5)))))
     36 x()
---> 37 torch.compile(x, backend="eager")()

21 frames
[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1500             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1501         else:
-> 1502             return self._call_impl(*args, **kwargs)
   1503 
   1504     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1509                 or _global_backward_pre_hooks or _global_backward_hooks
   1510                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1511             return forward_call(*args, **kwargs)
   1512         # Do not call functions when jit is used
   1513         full_backward_hooks, non_full_backward_hooks = [], []

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py](https://localhost:8080/#) in _fn(*args, **kwargs)
    293             dynamic_ctx.__enter__()
    294             try:
--> 295                 return fn(*args, **kwargs)
    296             finally:
    297                 set_eval_frame(prior)

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1500             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1501         else:
-> 1502             return self._call_impl(*args, **kwargs)
   1503 
   1504     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1509                 or _global_backward_pre_hooks or _global_backward_hooks
   1510                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1511             return forward_call(*args, **kwargs)
   1512         # Do not call functions when jit is used
   1513         full_backward_hooks, non_full_backward_hooks = [], []

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py](https://localhost:8080/#) in catch_errors(frame, cache_size, frame_state)
    446 
    447         with compile_lock:
--> 448             return callback(frame, cache_size, hooks, frame_state)
    449 
    450     catch_errors._torchdynamo_orig_callable = callback  # type: ignore[attr-defined]

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py](https://localhost:8080/#) in _convert_frame(frame, cache_size, hooks, frame_state)
    525         counters["frames"]["total"] += 1
    526         try:
--> 527             result = inner_convert(frame, cache_size, hooks, frame_state)
    528             counters["frames"]["ok"] += 1
    529             return result

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py](https://localhost:8080/#) in _fn(*args, **kwargs)
    125         cleanup = setup_compile_debug()
    126         try:
--> 127             return fn(*args, **kwargs)
    128         finally:
    129             cleanup.close()

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py](https://localhost:8080/#) in _convert_frame_assert(frame, cache_size, hooks, frame_state)
    358         )
    359 
--> 360         return _compile(
    361             frame.f_code,
    362             frame.f_globals,

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/utils.py](https://localhost:8080/#) in time_wrapper(*args, **kwargs)
    178             with torch.profiler.record_function(f"{key} (dynamo_timed)"):
    179                 t0 = time.time()
--> 180                 r = func(*args, **kwargs)
    181                 time_spent = time.time() - t0
    182             # print(f"Dynamo timer: key={key}, latency={latency:.2f} sec")

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py](https://localhost:8080/#) in _compile(code, globals, locals, builtins, compiler_fn, one_graph, export, export_constraints, hooks, frame, frame_state)
    428         for attempt in itertools.count():
    429             try:
--> 430                 out_code = transform_code_object(code, transform)
    431                 orig_code_map[out_code] = code
    432                 break

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/bytecode_transformation.py](https://localhost:8080/#) in transform_code_object(code, transformations, safe)
    998     propagate_line_nums(instructions)
    999 
-> 1000     transformations(instructions, code_options)
   1001     return clean_and_assemble_instructions(instructions, keys, code_options)[1]
   1002 

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/convert_frame.py](https://localhost:8080/#) in transform(instructions, code_options)
    413         )
    414         with tracing(tracer.output.tracing_context):
--> 415             tracer.run()
    416         output = tracer.output
    417         assert output is not None

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py](https://localhost:8080/#) in run(self)
   2023 
   2024     def run(self):
-> 2025         super().run()
   2026 
   2027     def match_nested_cell(self, name, cell):

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py](https://localhost:8080/#) in run(self)
    706                     self.instruction_pointer is not None
    707                     and not self.output.should_exit
--> 708                     and self.step()
    709                 ):
    710                     pass

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py](https://localhost:8080/#) in step(self)
    666                 self.f_code.co_filename, self.lineno, self.f_code.co_name
    667             )
--> 668             getattr(self, inst.opname)(inst)
    669 
    670             return inst.opname != "RETURN_VALUE"

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py](https://localhost:8080/#) in GET_ITER(self, inst)
   1092 
   1093     def GET_ITER(self, inst):
-> 1094         self.call_function(BuiltinVariable(iter), [self.pop()], {})
   1095 
   1096     @break_graph_if_unsupported(push=1)

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/symbolic_convert.py](https://localhost:8080/#) in call_function(self, fn, args, kwargs)
    557         ):
    558             raise AssertionError(f"Attempt to trace forbidden callable {inner_fn}")
--> 559         self.push(fn.call_function(self, args, kwargs))
    560 
    561     def update_locals_and_stack(self, oldvar: VariableTracker, newvar: VariableTracker):

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/builtin.py](https://localhost:8080/#) in call_function(self, tx, args, kwargs)
    575         if handler:
    576             try:
--> 577                 result = handler(tx, *args, **kwargs)
    578                 if result is not None:
    579                     return result.add_options(options)

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/builtin.py](https://localhost:8080/#) in _call_iter_tuple_list(self, tx, obj, *args, **kwargs)
    754                 mutable_local=MutableLocal(),
    755             )
--> 756         elif obj.has_unpack_var_sequence(tx):
    757             guards = set()
    758             if obj.source and not is_constant_source(obj.source):

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/base.py](https://localhost:8080/#) in has_unpack_var_sequence(self, tx)
    214     def has_unpack_var_sequence(self, tx):
    215         try:
--> 216             self.unpack_var_sequence(tx)
    217             return True
    218         except NotImplementedError:

[/usr/local/lib/python3.10/dist-packages/torch/_dynamo/variables/nn_module.py](https://localhost:8080/#) in unpack_var_sequence(self, tx)
     97             return result
     98 
---> 99         assert isinstance(
    100             base, (torch.nn.ModuleList, torch.nn.ParameterList, torch.nn.Sequential)
    101         ), typestr(base)

AssertionError: BufferList
```

### Minified repro

```
import torch
import torch.nn as nn

class BufferList(nn.Module):
    """
    Similar to nn.ParameterList, but for buffers
    """

    def __init__(self, buffers=None):
        super(BufferList, self).__init__()
        if buffers is not None:
            self.extend(buffers)

    def extend(self, buffers):
        offset = len(self)
        for i, buffer in enumerate(buffers):
            self.register_buffer(str(offset + i), buffer)
        return self

    def __len__(self):
        return len(self._buffers)

    def __iter__(self):
        return iter(self._buffers.values())

class Containing(nn.Module):
  def __init__(self, buflist):
    super().__init__()
    self.buflist = buflist
  def forward(self):
    for x in self.buflist:
      print(x.shape)

print(torch.__version__) #2.1.0.dev20230617+cu118
x = Containing(BufferList((torch.rand((5,5)), torch.rand((5,5)))))
x() # works
torch.compile(x, backend="eager")() # crashes
```

### Versions

Collecting environment information...
PyTorch version: 2.1.0.dev20230617+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.25.2
Libc version: glibc-2.31

Python version: 3.10.12 (main, Jun  7 2023, 12:45:35) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.107+-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: 11.8.89
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/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.0
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:                   46 bits physical, 48 bits virtual
CPU(s):                          2
On-line CPU(s) list:             0,1
Thread(s) per core:              2
Core(s) per socket:              1
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           79
Model name:                      Intel(R) Xeon(R) CPU @ 2.20GHz
Stepping:                        0
CPU MHz:                         2200.146
BogoMIPS:                        4400.29
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       32 KiB
L1i cache:                       32 KiB
L2 cache:                        256 KiB
L3 cache:                        55 MiB
NUMA node0 CPU(s):               0,1
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Vulnerable; SMT Host state unknown
Vulnerability Meltdown:          Vulnerable
Vulnerability Mmio stale data:   Vulnerable
Vulnerability Retbleed:          Vulnerable
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:        Vulnerable, IBPB: disabled, STIBP: disabled, PBRSB-eIBRS: Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Vulnerable
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 tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] pytorch-triton==2.1.0+440fd1bf20
[pip3] torch==2.1.0.dev20230617+cu118
[pip3] torchaudio==2.0.2+cu118
[pip3] torchdata==0.6.1
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.15.2
[pip3] torchvision==0.15.2+cu118
[pip3] triton==2.0.0
[conda] Could not collect

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305