Repository: pytorch
Issue: Cannot view a tensor with shape torch.Size([2, 1, 8, 32]) and strides (32, 512, 64, 1) as a tensor with shape (2, 256)! for transfomer MHA with permute, view
Labels: needs reproduction,module: onnx,oncall: pt2,oncall: export

### 🐛 Describe the bug

When exporting my transformer model with torch.onnx.export(... dynamo=True, ...), any sequence length >= 2 gives the error. Sequence length 1 works but does not support dynamic shapes in inference. When dynamic_shapes is off, I get the above error, and when dynamic_shapes is set and seq_len >=2, I get an error about guard lengths similar to #126127. I tried latest nightly (9 PM PST on 11/19) but didn't work. Also noticed this issue is similar to #136543 and #139508. Thanks in advance for the help!

[onnx_export_2024-11-19_21-08-53-216442_conversion.md](https://github.com/user-attachments/files/17825309/onnx_export_2024-11-19_21-08-53-216442_conversion.md)

**Repro**
```
import torch
import torch.nn as nn
import torch.nn.init as init
from torchvision.models import densenet121, DenseNet121_Weights
from torch import Tensor

class FlattenAndContiguous(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, x):
        return x.flatten(1, 2).contiguous()

class Permute(nn.Module):
    def __init__(self, *dims: int): # asterisk accepts arbitary amount of arguments
        super().__init__()
        self.dims = dims

    def forward(self, x):
        return x.permute(*self.dims).contiguous() # reorders the tuple
    
class PosEncode1D(nn.Module):
    def __init__(self, d_model, dropout_percent, max_len, PE_temp):
        super().__init__()

        position = torch.arange(max_len).unsqueeze(1) # creates a vector (max_len x 1), 1 is needed for matmul operations
        dim_t = torch.arange(0, d_model, 2) # 2i term in the denominator exponent
        scaling = PE_temp **(dim_t/d_model) # entire denominator 

        pe = torch.zeros(max_len, d_model) # 
        pe[:, 0::2] = torch.sin(position / scaling) # every second term starting from 0 (even)
        pe[:, 1::2] = torch.cos(position / scaling) # every second term starting from 1 (odd)

        self.dropout = nn.Dropout(dropout_percent)
        self.register_buffer("pe", pe) # stores pe tensor to be used but not updated

    def forward(self, x):
        batch, sequence_length, d_model = x.shape         
        return self.dropout(x + self.pe[None, :sequence_length, :]) # None to broadcast across batch, adds element-wise [x + pe, . . .]


class PosEncode2D(nn.Module):
    def __init__(self, d_model, dropout_percent, max_len, PE_temp):
        super().__init__()
        # 1D encoding
        position = torch.arange(max_len).unsqueeze(1) 
        dim_t = torch.arange(0, d_model, 2)
        scaling = PE_temp **(dim_t/d_model)
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position / scaling)
        pe[:, 1::2] = torch.cos(position / scaling)

        pe_2D = torch.zeros(max_len, max_len, d_model) # some outer product magic
        for i in range(d_model):
            pe_2D[:, :, i] = pe[:, i].unsqueeze(1) + pe[:, i].unsqueeze(0)  # first unsqueeze changed from -1

        self.dropout = nn.Dropout(dropout_percent)
        self.register_buffer("pe", pe_2D) 

    def forward(self, x):
        batch, height, width, d_model = x.shape
        return self.dropout(x + self.pe[None, :height, :width, :]).contiguous()
    

class Model_1(nn.Module): # from https://actamachina.com/handwritten-mathematical-expression-recognition, CNN encoder and then transformer decoder
    def __init__(self, vocab_size, d_model, nhead, dim_FF, dropout, num_layers):
        super(Model_1, self).__init__()
        densenet = densenet121(weights=DenseNet121_Weights.DEFAULT) 
        
        self.encoder = nn.Sequential(
            nn.Sequential(*list(densenet.children())[:-1]), # remove the final layer, output (B, 1024, 12, 16)
            nn.Conv2d(1024, d_model, kernel_size=1), # 1x1 convolution, output of (B, d_model, W, H) ex. (1, 256, 12, 16)
            Permute(0, 3, 2, 1), 
            PosEncode2D(d_model=d_model, dropout_percent=dropout, max_len=150, PE_temp=10000), # output (1, 16, 12, 256)
            FlattenAndContiguous()
        )
    
        self.tgt_embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model, padding_idx=0) # simple lookup table, input indices 
        self.word_PE = PosEncode1D(d_model, dropout, max_len=150, PE_temp=10000)
        self.transformer_decoder = nn.TransformerDecoder( 
            nn.TransformerDecoderLayer(d_model, nhead, dim_FF, dropout, batch_first=True), # batch_first -> (batch, sequence, feature)
            num_layers,
        ) # input target and memory (last sequence of the encoder), then tgt_mask, memory_mask
        self.fc_out = nn.Linear(256, 217) # y = xA^T + b, distribution over all tokens in vocabulary
        self.d_model = Tensor(d_model).to(device)

        self._initialize_weights()  

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                init.xavier_uniform_(m.weight)  # Glorot initialization for linear layers
                if m.bias is not None:
                    init.zeros_(m.bias)  # Bias initialized to zeros
            elif isinstance(m, nn.Conv2d):
                init.xavier_uniform_(m.weight)  # Glorot initialization for conv layers
                if m.bias is not None:
                    init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding):
                init.xavier_uniform_(m.weight)  # Glorot initialization for embedding layers
    

    def decoder(self, features, tgt, tgt_mask):
        features = features.contiguous()
        tgt = tgt.contiguous()
        padding_mask = tgt.eq(0) # checks where elements of tgt are equal to zero
        tgt = self.tgt_embedding(tgt) * torch.sqrt(self.d_model) # tgt indices become embedding vectors and are scaled by sqrt of model size for stability
        tgt = self.word_PE(tgt) # adds positional encoding, size (B, seq_len, d_model)
        tgt = self.transformer_decoder(tgt=tgt, memory=features, tgt_mask=tgt_mask.to(torch.float32), tgt_key_padding_mask=padding_mask.to(torch.float32), tgt_is_causal=True) # type match
        #tgt = self.transformer_decoder(tgt=tgt, memory=features, tgt_mask=tgt_mask, tgt_is_causal=True) # type match
        output = self.fc_out(tgt) # size (B, seq_len, vocab_size
        return output
    
    def forward(self, src, tgt, tgt_mask):
        features = self.encoder(src)
        output = self.decoder(features, tgt, tgt_mask)
        return output
    


import torch.onnx
import onnxruntime as ort
import numpy as np

device = 'cuda'

full_model = Model_1(vocab_size=217, d_model=256, nhead=8, dim_FF=1024, dropout=0, num_layers=3)
full_model.to(device)
#full_model.load_state_dict(torch.load('runs/Exp8E8End_Acc=0.6702922582626343.pt', map_location = device, weights_only=True))
full_model.eval()

with torch.no_grad():
    seq_len_tgt = torch.export.Dim("seq_len_tgt", max=146)
    seq_len_0 = torch.export.Dim("seq_len_0", max=146)
    seq_len_1 = torch.export.Dim("seq_len_1", max=146)
    
    torch.onnx.export(
        full_model.to(device), 
        (torch.randn(1,3,384,512).to(device),
        torch.ones([1, 2], dtype=torch.long).to(device),
        torch.triu(torch.ones(2, 2) * float("-inf"), diagonal=1).to(device)
        ),
        "model_3.onnx",
        export_params = True,
        verbose = True,
        input_names = ["src", "tgt", "tgt_mask"],
        output_names = ["output1"],
        dynamo = True,
        #dynamic_shapes = {"src": {}, "tgt": {1: seq_len_tgt}, "tgt_mask": {0: seq_len_0, 1: seq_len_1}},
        external_data = True,
        report = True,
        optimize = True,
        verify = True, 
    )
```

### Versions

Collecting environment information...
PyTorch version: 2.6.0.dev20241112+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.1 LTS (x86_64)
GCC version: (Ubuntu 13.2.0-23ubuntu4) 13.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.39

Python version: 3.12.3 (main, Sep 11 2024, 14:17:37) [GCC 13.2.0] (64-bit runtime)
Python platform: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4070 Laptop GPU
Nvidia driver version: 560.94
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:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             22
On-line CPU(s) list:                0-21
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Core(TM) Ultra 9 185H
CPU family:                         6
Model:                              170
Thread(s) per core:                 2
Core(s) per socket:                 11
Socket(s):                          1
Stepping:                           4
BogoMIPS:                           6144.00
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 tsc_reliable nonstop_tsc cpuid pni pclmulqdq vmx 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 ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          528 KiB (11 instances)
L1i cache:                          704 KiB (11 instances)
L2 cache:                           22 MiB (11 instances)
L3 cache:                           24 MiB (1 instance)
Vulnerability Gather data sampling: Not affected
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:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: 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

Versions of relevant libraries:
[pip3] numpy==2.1.3
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] onnx==1.17.0
[pip3] onnxruntime==1.20.0
[pip3] onnxscript==0.1.0.dev20241112
[pip3] pytorch-triton==3.1.0+cf34004b8a
[pip3] torch==2.6.0.dev20241112+cu124
[pip3] torchaudio==2.5.0.dev20241112+cu124
[pip3] torchvision==0.20.0.dev20241112+cu124
[pip3] triton==3.1.0
[conda] Could not collect

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4