Repository: pytorch
Issue: An additional dimension appears in the return when using torch.compile on torch.nn.LSTM models
Labels: triaged,oncall: pt2,module: dynamo

### 🐛 Describe the bug

There is an additional dimension appearing in the second return value of a `torch.nn.LSTM` layer when we apply `torch.compile` (in all backends) to it. The additional dimension happens to both tensors wrapped by a tuple in the second return value. After some investigations, this issue only happens when all of the following conditions hold:
1. Apply `torch.compile` to models containing a `torch.nn.LSTM` layer (`torch.nn.RNN` and `torch.nn.GRU` are fine).
2. The tensors and models are on device (GPU).
3. When we view the "intermediate" shape `torch.nn.LSTM`'s output.

To be more clear, running my reproducer below will get this output, where the additional dimension appears after applying `torch.compile` (the intermediate shape is changed from "[2, 8, 640]" to "[2, 1, 8, 640]"):
```
====================== Eager PyTorch ======================
Return shape of hid[0]:
torch.Size([2, 8, 640])

Intermediate shape of hid[0]:
tensor([2.]) tensor([8.]) tensor([640.]) tensor([0.])

====================== Torch Compile ======================
Return shape of hid[0]:
torch.Size([2, 8, 640])

Intermediate shape of hid[0]:
tensor([2.]) tensor([1.]) tensor([8.]) tensor([640.])
```

I have also used `TORCH_LOGS="+dynamo,bytecode"` for debugging, and its shortened output is provided in "Error logs". The incorrect intermediate shape can be found by searching "TRACED GRAPH" in "Error logs".

Another tricky part here is that although the intermediate shape is wrong, the shape of the return value is correct. However, the Python bytecode shown in the "MODIFIED BYTECODE" section in "Error logs" does not contain operations which may alter the tensors' dimension.

### Error logs

```
V0716 00:02:07.514000 139634126092096 torch/_dynamo/output_graph.py:976] [0/0] COMPILING GRAPH due to GraphCompileReason(reason='return_value', user_stack=[<FrameSummary file /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py, line 31 in forward>], graph_break=False)
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code] TRACED GRAPH
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]  ===== __compiled_fn_1 =====
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]  /media/data/pytorch/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]     def forward(self, L_x_: "f32[1, 8, 640][5120, 640, 1]cuda:0", L_h_0_: "f32[2, 8, 640][5120, 640, 1]cuda:0", L_h_1_: "f32[2, 8, 640][5120, 640, 1]cuda:0", G_d1_: "f32[1][1]cpu", G_d2_: "f32[1][1]cpu", G_d3_: "f32[1][1]cpu", G_d4_: "f32[1][1]cpu"):
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         l_x_ = L_x_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         l_h_0_ = L_h_0_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         l_h_1_ = L_h_1_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         g_d1_ = G_d1_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         g_d2_ = G_d2_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         g_d3_ = G_d3_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         g_d4_ = G_d4_
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]          # File: /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py:22 in forward, code: g, hid = self.lstm(x, h)
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         l__self___lstm = self.L__self___lstm(l_x_, (l_h_0_, l_h_1_));  l_x_ = l_h_0_ = l_h_1_ = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         g: "f32[1, 8, 640][5120, 640, 1]cuda:0" = l__self___lstm[0]
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         getitem_1 = l__self___lstm[1];  l__self___lstm = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         getitem_2: "f32[2, 1, 8, 640][5120, 5120, 640, 1]cuda:0" = getitem_1[0]
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         getitem_3: "f32[2, 1, 8, 640][5120, 5120, 640, 1]cuda:0" = getitem_1[1];  getitem_1 = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]          # File: /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py:25 in forward, code: d1.copy_(hid[0].shape[0])
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         copy_: "f32[1][1]cpu" = g_d1_.copy_(2);  g_d1_ = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]          # File: /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py:26 in forward, code: d2.copy_(hid[0].shape[1])
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         copy__1: "f32[1][1]cpu" = g_d2_.copy_(1);  g_d2_ = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]          # File: /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py:27 in forward, code: d3.copy_(hid[0].shape[2])
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         copy__2: "f32[1][1]cpu" = g_d3_.copy_(8);  g_d3_ = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]          # File: /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py:29 in forward, code: d4.copy_(hid[0].shape[3])
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         copy__3: "f32[1][1]cpu" = g_d4_.copy_(640);  g_d4_ = None
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         return (g, getitem_2, getitem_3)
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code]         
V0716 00:02:07.516000 139634126092096 torch/_dynamo/output_graph.py:1295] [0/0] [__graph_code] 
I0716 00:02:07.517000 139634126092096 torch/_dynamo/logging.py:56] [0/0] Step 2: calling compiler function eager
I0716 00:02:07.517000 139634126092096 torch/_dynamo/logging.py:56] [0/0] Step 2: done compiler function eager
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode] MODIFIED BYTECODE forward /media/data/test/repro/greedy_decoding/lstm_squeeze/lstm_repro.py line 21 
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]  21           0 LOAD_GLOBAL              8 (__compiled_fn_1)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]               2 LOAD_FAST                1 (x)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]               4 LOAD_FAST                2 (h)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]               6 LOAD_CONST               1 (0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]               8 BINARY_SUBSCR
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              10 LOAD_FAST                2 (h)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              12 LOAD_CONST               2 (1)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              14 BINARY_SUBSCR
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              16 LOAD_GLOBAL              1 (d1)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              18 LOAD_GLOBAL              4 (d2)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              20 LOAD_GLOBAL              5 (d3)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              22 LOAD_GLOBAL              7 (d4)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              24 CALL_FUNCTION            7
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              26 STORE_FAST               5 (graph_out_0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              28 LOAD_FAST                5 (graph_out_0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              30 LOAD_CONST               1 (0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              32 BINARY_SUBSCR
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              34 LOAD_FAST                5 (graph_out_0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              36 LOAD_CONST               2 (1)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              38 BINARY_SUBSCR
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              40 LOAD_FAST                5 (graph_out_0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              42 LOAD_CONST               3 (2)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              44 BINARY_SUBSCR
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              46 BUILD_TUPLE              2
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              48 BUILD_TUPLE              2
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              50 DELETE_FAST              5 (graph_out_0)
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode]              52 RETURN_VALUE
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode] 
V0716 00:02:07.519000 139634126092096 torch/_dynamo/convert_frame.py:640] [0/0] [__bytecode] 
```

### Minified repro

```
import torch
# enable RNN in torchdynamo
from torch._dynamo import config
config.allow_rnn = True

device = "cuda"
x = torch.zeros([1, 8, 640], device=device)
h = (torch.zeros([2, 8, 640], device=device), torch.zeros([2, 8, 640], device=device))
d1 = torch.zeros(1, device="cpu", pin_memory=True)
d2 = torch.zeros(1, device="cpu", pin_memory=True)
d3 = torch.zeros(1, device="cpu", pin_memory=True)
d4 = torch.zeros(1, device="cpu", pin_memory=True)

class MY_LSTM(torch.nn.Module):
    def __init__(self):
        super(MY_LSTM, self).__init__()
        self.lstm = torch.nn.LSTM(
            input_size=640, hidden_size=640, num_layers=2, dropout=0.2, proj_size=0
        )

    def forward(self, x, h):
        g, hid = self.lstm(x, h)

        # copy the "intermediate shape" of hid
        d1.copy_(hid[0].shape[0])
        d2.copy_(hid[0].shape[1])
        d3.copy_(hid[0].shape[2])
        if hid[0].dim() > 3:
            d4.copy_(hid[0].shape[3])

        return g, hid

lstm = MY_LSTM().to(device=device)
compiled_lstm = torch.compile(lstm, backend="eager")

print("\n====================== Eager PyTorch ======================")
_, hid_ret = lstm(x, h)

print("Return shape of hid[0]:")
print(hid_ret[0].shape)
print("\nIntermediate shape of hid[0]:")
print(d1, d2, d3, d4)

print("\n====================== Torch Compile ======================")
_, hid_ret = compiled_lstm(x, h)

print("Return shape of hid[0]:")
print(hid_ret[0].shape)
print("\nIntermediate shape of hid[0]:")
print(d1, d2, d3, d4)
```

### Versions

```
Collecting environment information...
PyTorch version: 2.5.0a0+git9ab0b9f
Is debug build: True
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.29.3
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jul  5 2023, 18:54:27) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L40S
Nvidia driver version: 550.54.15
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:                   45 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          25
On-line CPU(s) list:             0-24
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 9454 48-Core Processor
CPU family:                      25
Model:                           17
Thread(s) per core:              1
Core(s) per socket:              25
Socket(s):                       1
Stepping:                        1
BogoMIPS:                        5491.74
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 tsc_reliable nonstop_tsc cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext invpcid_single ssbd ibrs ibpb vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512_bf16 clzero wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid overflow_recov succor fsrm flush_l1d
Hypervisor vendor:               VMware
Virtualization type:             full
L1d cache:                       800 KiB (25 instances)
L1i cache:                       800 KiB (25 instances)
L2 cache:                        25 MiB (25 instances)
L3 cache:                        32 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-24
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; Retpolines, IBPB conditional, IBRS_FW, STIBP disabled, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] onnx==1.16.0
[pip3] open-clip-torch==2.24.0
[pip3] pytorch-lightning==2.2.4
[pip3] torch==2.5.0a0+gitfca6212
[pip3] torchdiffeq==0.2.3
[pip3] torchmetrics==1.4.0
[pip3] torchsde==0.2.6
[pip3] triton==2.3.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.8           py310h5eee18b_0  
[conda] mkl_random                1.2.4           py310hdb19cb5_0  
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpy-base                1.24.3          py310hb5e798b_1  
[conda] open-clip-torch           2.24.0                   pypi_0    pypi
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-lightning         2.2.4                    pypi_0    pypi
[conda] torch                     2.5.0a0+gitfca6212           dev_0    <develop>
[conda] torchdiffeq               0.2.3                    pypi_0    pypi
[conda] torchmetrics              1.4.0                    pypi_0    pypi
[conda] torchsde                  0.2.6                    pypi_0    pypi
[conda] torchtriton               2.3.0                     py310    pytorch
```

cc @ezyang @anijain2305 @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames