Repository: pytorch
Issue: Batch size can't be varied if you export self attention to onnx
Labels: module: onnx,module: windows,triaged

### 🐛 Describe the bug

I've been making a lot of tickets about differentiating through things lately, but here's a much more basic one: if I export a self attention layer from a transformer to on onnx, without even running any backprop stuff, the onnx appears to be broken if I try to run it with different batch sizes and sequence lengths. Execute this code to reproduce the issue:

```python
import torch
import onnxruntime as ort


class Model(torch.nn.Module):

    def __init__(self):
        super().__init__()

        self.latent_dim = 256
        self.num_heads = 4
        self.ff_size=1024
        self.dropout=0.1
        self.activation="gelu"
        self.num_layers = 4

        self.root_seqTransEncoderLayer = torch.nn.TransformerEncoderLayer(d_model=self.latent_dim,
                                                               nhead=self.num_heads,
                                                               dim_feedforward=self.ff_size,
                                                               dropout=self.dropout,
                                                               activation=self.activation)
        
    def forward(self, inputs):
        x = inputs[0]
        x = self.root_seqTransEncoderLayer.self_attn(x, x, x,
                           attn_mask=None,
                           key_padding_mask=None,
                           need_weights=False, is_causal=False)[0]
        
        return x


batch_size=20
seq_len=2

mdl = Model()
for p in mdl.parameters():
    p.requires_grad_(False)


print("export model")
torch.onnx.export(
    Model(),
    [torch.randn([batch_size, seq_len, 256])],
    "modelthing.onnx",
    dynamic_axes={'xseq': {0: 'batch_size', 1:'seq_len'}},
    input_names=["xseq"],
    opset_version=17,
    output_names=["lossgrad"],
    verbose=True
)

# load and evaluate with a different shape:
ort_session = ort.InferenceSession("modelthing.onnx")

print("run onnx with export shape")
ort_session.run(None, { "xseq":torch.randn([batch_size, seq_len, 256]).numpy() })
print("success")

seq_len=5
# comment this out and it should work:
batch_size=10

print("run onnx with different batch size")
ort_session.run(None, { "xseq":torch.randn([batch_size, seq_len, 256]).numpy() })
print("success")
```

Here's the error I get:

```
2024-03-20 10:46:07.3028954 [E:onnxruntime:, sequential_executor.cc:514 onnxruntime::ExecuteKernel] Non-zero status code returned while running Reshape node. Name:'/self_attn/Reshape_4' Status Message: C:\a\_work\1\s\onnxruntime\core\providers\cpu\tensor\reshape_helper.h:41 onnxruntime::ReshapeHelper::ReshapeHelper gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{10,5,256}, requested shape:{20,20,64}

Traceback (most recent call last):
  File "C:\Git\physdiff\dynamicshapesbroken.py", line 68, in <module>
    ort_session.run(None, { "xseq":torch.randn([batch_size, seq_len, 256]).numpy() })
  File "C:\Users\dminor\AppData\Roaming\Python\Python310\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py", line 217, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'/self_attn/Reshape_4' Status Message: C:\a\_work\1\s\onnxruntime\core\providers\cpu\tensor\reshape_helper.h:41 onnxruntime::ReshapeHelper::ReshapeHelper gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{10,5,256}, requested shape:{20,20,64}
```

I don't know if this works better with the dynamo export, but that's not currently an option for me as I can mostly differentiate through things with torch.onnx.export() now and it's completely unsupported in dynamo.


### Versions

PyTorch version: 2.2.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Enterprise
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.25.1
Libc version: N/A

Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr 5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 10.0.130
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090 Ti
Nvidia driver version: 536.23
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3501
DeviceID=CPU0
Family=107
L2CacheSize=16384
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=3501
Name=AMD Ryzen Threadripper PRO 3975WX 32-Cores
ProcessorType=3
Revision=12544

Versions of relevant libraries:
[pip3] functorch==2.0.0
[pip3] lovely-numpy==0.2.8
[pip3] numpy==1.24.3
[pip3] onnx==1.14.1
[pip3] onnx-graphsurgeon==0.3.27
[pip3] onnxconverter-common==1.13.0
[pip3] onnxruntime==1.15.1
[pip3] optree==0.10.0
[pip3] pytorch-lightning==1.4.2
[pip3] tf2onnx==1.16.0
[pip3] torch==2.2.1+cu118
[pip3] torch-cluster==1.6.1
[pip3] torch-fidelity==0.3.0
[pip3] torch-geometric==2.3.0
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torch-spline-conv==1.2.2
[pip3] torchaudio==2.2.1+cu118
[pip3] torchmetrics==0.6.0
[pip3] torchvision==0.17.1+cu118
[conda] blas 1.0 mkl
[conda] cudatoolkit 11.6.0 hc0ea762_10 conda-forge
[conda] libblas 3.9.0 16_win64_mkl conda-forge
[conda] libcblas 3.9.0 16_win64_mkl conda-forge
[conda] liblapack 3.9.0 16_win64_mkl conda-forge
[conda] mkl 2022.1.0 h6a75c08_874 conda-forge
[conda] mkl-include 2023.2.0 intel_49496 intel
[conda] mkl-static 2023.2.0 intel_49496 intel
[conda] numpy 1.24.2 py310hd02465a_0 conda-forge
[conda] pytorch 1.12.0 py3.10_cuda11.6_cudnn8_0 pytorch
[conda] pytorch-mutex 1.0 cuda pytorch

cc @peterjc123 @mszhanyi @skyline75489 @nbcsm @vladimir-aubrecht @iremyux @Blackhex @cristianPanaite