[Bug][Relax][Frontend][Torch] `torch.einsum` with repeated subscripts (diagonal/trace, e.g. `'ii->i'`, `'ii->'`, `'...ii->...i'`) fails to convert: `aten.diagonal` unsupported

[siyiweigeHEW]

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:

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

"""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