LiSsHhUuAaIi opened a new issue, #18441:
URL: https://github.com/apache/tvm/issues/18441
### Description
When converting a PyTorch model containing
`torch.nn.functional.scaled_dot_product_attention` with 2D input tensors to TVM
Relax module via `torch.export`, an InternalError occurs during the conversion
process. The TVM frontend incorrectly assumes 4D input dimensions for the
attention operation when the actual input is 2D.
### Expected behavior
The PyTorch model with `scaled_dot_product_attention` should be successfully
converted to TVM Relax module regardless of input dimensionality, as it runs
correctly in native PyTorch.
### Actual behavior
An InternalError occurs during `from_exported_program` conversion with the
message `PermuteDims expects the number of input axes to equal the ndim of the
input tensor. However, the tensor ndim is 2 while the given number of axes is
4`, indicating a dimension mismatch in the PermuteDims operation.
### Environment
* **OS:** Ubuntu 20.04.6 LTS
* **TVM version:** 0.23.dev0
* **Python version:** 3.11.14
### Steps to reproduce
```python
import torch
import torch.nn as nn
import tvm
from tvm import relax
class MinimalAttentionModel(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return torch.nn.functional.scaled_dot_product_attention(x, x, x,
is_causal=False)
model = MinimalAttentionModel()
model.eval()
# 2D input triggers the bug
x = torch.randn(8, 32)
# PyTorch execution works
with torch.no_grad():
output = model(x)
# PyTorch export works
exported_program = torch.export.export(model, (x,))
# TVM conversion fails
from tvm.relax.frontend.torch import from_exported_program
mod = from_exported_program(exported_program) # InternalError here
```
### Error Log
```
Traceback (most recent call last):
File "test.py", line 29, in <module>
mod = from_exported_program(exported_program) # InternalError here
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
....
tvm.error.InternalError: PermuteDims expects the number of input axes to
equal the ndim of the input tensor. However, the tensor ndim is 2 while the
given number of axes is 4
```
### Triage
* needs-triage
* bug
* frontend: pytorch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
