Repository: pytorch
Issue: [Torch.fx] Torch fx failed to trace torch extension library
Labels: triaged,module: fx

### 🐛 Describe the bug

Hi,

I'm currently focusing on point cloud detection workloads, and I rely heavily on the torch extension libraries like mmdet3d. I recently attempted to trace a straightforward torch cpp extension using torch.fx, but unfortunately, I encountered an error.

For reference, the code I'm working with is directly from the official PyTorch documentation, specifically this [tutorial on cpp extensions](https://pytorch.org/tutorials/advanced/cpp_extension.html).

Is there any known limitation or workaround when it comes to using torch.fx with torch cpp extensions? I'd appreciate any guidance or recommendations on this matter.
``` python
class LLTM(torch.nn.Module):
    def __init__(self, input_features, state_size):
        super(LLTM, self).__init__()
        self.input_features = input_features
        self.state_size = state_size
        self.weights = torch.nn.Parameter(
            torch.empty(3 * state_size, input_features + state_size))
        self.bias = torch.nn.Parameter(torch.empty(3 * state_size))
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.state_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, +stdv)

    def forward(self, input, state):
        return LLTMFunction.apply(input, self.weights, self.bias, *state)
    
```
Here is the tracing code.
``` python
import time

import torch

batch_size = 16
input_features = 32
state_size = 128

X = torch.randn(batch_size, input_features)
h = torch.randn(batch_size, state_size)
C = torch.randn(batch_size, state_size)

rnn = LLTM(input_features, state_size)
compile_rnn = torch.compile(rnn, backend="inductor")
new_h, new_C = compile_rnn(X, (h, C))

import torch.fx

tracer = torch.fx.symbolic_trace(rnn)
print(tracer)
```
Then I got these errors:
```
Traceback (most recent call last):
  File "/home/leosys/Documents/test_fx/lltm.py", line 62, in <module>
    tracer = torch.fx.symbolic_trace(rnn)
  File "/home/leosys/.conda/envs/gencom/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 1150, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/leosys/.conda/envs/gencom/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/home/leosys/.conda/envs/gencom/lib/python3.9/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/home/leosys/.conda/envs/gencom/lib/python3.9/site-packages/torch/fx/_symbolic_trace.py", line 817, in trace
    (self.create_arg(fn(*args)),),
  File "/home/leosys/Documents/test_fx/lltm.py", line 41, in forward
    return LLTMFunction.apply(input, self.weights, self.bias, *state)
  File "/home/leosys/.conda/envs/gencom/lib/python3.9/site-packages/torch/fx/proxy.py", line 409, in __iter__
    return self.tracer.iter(self)
  File "/home/leosys/.conda/envs/gencom/lib/python3.9/site-packages/torch/fx/proxy.py", line 309, in iter
    raise TraceError('Proxy object cannot be iterated. This can be '
torch.fx.proxy.TraceError: Proxy object cannot be iterated. This can be attempted when the Proxy is used in a loop or as a *args or **kwargs function argument. See the torch.fx docs on pytorch.org for a more detailed explanation of what types of control flow can be traced, and check out the Proxy docstring for help troubleshooting Proxy iteration errors
```

### Versions

PyTorch version: 2.1.0.dev20230821
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.35

Python version: 3.9.17 (main, Jul  5 2023, 20:41:20)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.2.0-26-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080
Nvidia driver version: 525.125.06
cuDNN version: Could not collect
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
Address sizes:                   48 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          16
On-line CPU(s) list:             0-15
Vendor ID:                       AuthenticAMD
Model name:                      AMD Ryzen 7 5800X 8-Core Processor
CPU family:                      25
Model:                           33
Thread(s) per core:              2
Core(s) per socket:              8
Socket(s):                       1
Stepping:                        2
Frequency boost:                 enabled
CPU max MHz:                     4850.1948
CPU min MHz:                     2200.0000
BogoMIPS:                        7585.37
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm
Virtualization:                  AMD-V
L1d cache:                       256 KiB (8 instances)
L1i cache:                       256 KiB (8 instances)
L2 cache:                        4 MiB (8 instances)
L3 cache:                        32 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-15
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
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] pytorch-triton==2.1.0+e6216047b8
[pip3] torch==2.1.0.dev20230820+cu118
[pip3] torch-geometric==2.3.1
[pip3] torchaudio==2.1.0.dev20230821+cu118
[pip3] torchvision==0.16.0.dev20230821+cu118
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] libjpeg-turbo             2.0.0                h9bf148f_0    pytorch-nightly
[conda] mkl                       2023.1.0         h213fc3f_46343  
[conda] mkl-service               2.4.0            py39h5eee18b_1  
[conda] mkl_fft                   1.3.6            py39h417a72b_1  
[conda] mkl_random                1.2.2            py39h417a72b_1  
[conda] numpy                     1.24.1                   pypi_0    pypi
[conda] numpy-base                1.25.2           py39hb5e798b_0  
[conda] pytorch                   2.1.0.dev20230821 py3.9_cuda11.8_cudnn8.7.0_0    pytorch-nightly
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch-nightly
[conda] pytorch-mutex             1.0                        cuda    pytorch-nightly
[conda] pytorch-triton            2.1.0+e6216047b8          pypi_0    pypi
[conda] torch                     2.1.0.dev20230820+cu118          pypi_0    pypi
[conda] torch-geometric           2.3.1                    pypi_0    pypi
[conda] torchaudio                2.1.0.dev20230821+cu118          pypi_0    pypi
[conda] torchtriton               2.1.0+e6216047b8            py39    pytorch-nightly
[conda] torchvision               0.16.0.dev20230821+cu118          pypi_0    pypi

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv