siyiweigeHEW opened a new issue, #20228:
URL: https://github.com/apache/tvm/issues/20228
### Expected behavior
A valid PyTorch model that calls `torch.einsum` with a repeated subscript on
the left-hand side (diagonal extraction / trace) should convert through
`tvm.relax.frontend.torch.from_exported_program`. These are standard,
documented einsum usages and run fine in PyTorch:
```python
import torch
x = torch.randn(3, 3)
torch.einsum("ii->i", x) # diagonal -> (3,)
torch.einsum("ii->", x) # trace -> scalar
torch.einsum("...ii->...i", x.unsqueeze(0).expand(3, 3, 3)) # batched
diagonal
```
### Actual behavior
`from_exported_program` rejects every one of these with:
```
AssertionError: Unsupported function types ['diagonal.default']
```
Root cause: `from_exported_program` runs
`exported_program.run_decompositions()` by default
(`python/tvm/relax/frontend/torch/exported_program_translator.py:1841-1845`).
PyTorch's decomposition lowers repeated-subscript einsum to `aten.diagonal` +
`permute` (+ `sum` for trace):
```
input: torch.einsum('ii->i', x)
after run_decompositions():
aten.diagonal.default
aten.permute.default
aten.permute.default
```
`aten.diagonal.default` is missing from the torch frontend `convert_map`, so
`_check_unsupported_func_type`
(`python/tvm/relax/frontend/torch/base_fx_graph_translator.py:186-194`)
asserts. The same error also hits the direct ops `torch.diagonal` and
`torch.trace` (which decomposes to `aten.diagonal` + clone + sum).
As evidence that the defect is specifically the missing `aten.diagonal`
handling (and **not** `relax.op.einsum` semantics), the same models convert and
run correctly when decomposition is skipped (`from_exported_program(ep,
run_ep_decomposition=False)`, which routes to the frontend's `_einsum` →
`relax.op.einsum`), with `max|diff| = 0` vs PyTorch. All *regular* einsum
patterns (matmul, transpose, dot, outer, batch matmul, ellipsis
broadcasting/summation, implicit output, 3-operand) also match PyTorch exactly
on the default path.
Verified failing equations (all valid in torch, all rejected by TVM):
`'ii->i'`, `'ii->'`, `'...ii->...i'`, `'...ii->...'`, `'iij->ij'`,
`'iji->j'`, `'ijj->i'`, `'abca,abcb->c'` (attention-style), `'iij,kk->ij'`,
`'...iij->ij'`.
### Environment
- OS: Linux
- TVM: v0.24.dev0-46-g262c6d2e04 (baseline of this report; also verified
against current `main` on 2026-08-26 — `convert_map` still has no
`diagonal.default`)
- Python: 3.11
- torch: 2.10.0+cu128
### Steps to reproduce
```python
"""Repro: torch.einsum with repeated subscripts (diagonal/trace) fails to
convert."""
import torch
import torch.nn as nn
from tvm.relax.frontend.torch import from_exported_program
class DiagEinsum(nn.Module):
def forward(self, x):
return torch.einsum("ii->i", x)
class TraceEinsum(nn.Module):
def forward(self, x):
return torch.einsum("ii->", x)
x = torch.randn(3, 3)
for label, module in [("diag('ii->i')", DiagEinsum()), ("trace('ii->')",
TraceEinsum())]:
print(f"--- {label} ---")
print(" torch output shape:", tuple(module(x).shape))
ep = torch.export.export(module.eval(), (x,))
print(" run_decompositions ops:",
[str(n.target) for n in ep.run_decompositions().graph.nodes if
n.op == "call_function"])
try:
from_exported_program(ep)
print(" TVM: OK")
except Exception as e:
print(f" TVM: {type(e).__name__}: {e}")
print()
```
Actual output:
```
--- diag('ii->i') ---
torch output shape: (3,)
run_decompositions ops: ['aten.diagonal.default', 'aten.permute.default',
'aten.permute.default']
TVM: AssertionError: Unsupported function types ['diagonal.default']
--- trace('ii->') ---
torch output shape: ()
run_decompositions ops: ['aten.diagonal.default', 'aten.permute.default',
'aten.permute.default', 'aten.sum.dim_IntList']
TVM: AssertionError: Unsupported function types ['diagonal.default']
```
### Suggested fix
Add an `aten.diagonal` converter (and/or handle the `diagonal` op produced
by einsum decomposition) in the torch frontend `convert_map`, so
repeated-subscript einsum converts. `torch.diagonal(x, offset, dim1, dim2)` can
be expressed with `relax.op.take`/slicing on the diagonal, or the einsum nodes
can be kept intact instead of decomposed. This is the same root cause for
`torch.diagonal` / `torch.trace` not converting.
### Triage
* needs-triage
* bug
* relax
* frontend/torch
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
