[Fix][Relax][Frontend][Torch] Support aten.diagonal from decomposed repeated-subscript einsum

[siyiweigeHEW]

Fixes: #20228

Summary

from_exported_program runs exported_program.run_decompositions() by
default, and PyTorch's decomposition lowers torch.einsum with repeated
subscripts (diagonal / trace, e.g. "ii->i", "ii->", "...ii->...i") to
aten.diagonal + permute (+ sum for the trace). aten.diagonal.default
was missing from the torch frontend convert_map, so every such valid
model failed with:

AssertionError: Unsupported function types ['diagonal.default']

This PR adds an aten.diagonal converter and registers it in both the
exported-program and from_fx convert maps, so repeated-subscript einsum —
and the directly-affected ops torch.diagonal / torch.trace — convert and
run. Verified failing equations from the issue all convert with max|diff| = 0
vs PyTorch.

Root cause

BaseFXGraphImporter._check_unsupported_func_type asserts when a
call_function node's target is not in convert_map. For the einsum family
above, run_decompositions introduces aten.diagonal.default nodes that the
torch frontend had no handler for, so conversion aborts at the assertion. This
is the same root cause for the direct ops torch.diagonal (lowered to
diagonal.default as-is) and torch.trace (lowered to diagonal + clone +
sum). Skipping decomposition (run_ep_decomposition=False) keeps the einsum
node intact and works — confirming the defect is the missing diagonal
handling, not relax.op.einsum semantics.

Fix

Add BaseFXGraphImporter._diagonal in base_fx_graph_translator.py, lowering
diagonal(input, offset=0, dim1=0, dim2=1) as:

1. relax.op.permute_dims — move dim1 / dim2 to the trailing two axes;
2. two relax.op.strided_slice — crop each trailing axis to the diagonal
   length min(extent1, extent2 ± offset) (offset-adjusted), so the two
   trailing extents are equal;
3. relax.op.einsum([x], "...zz->...z") — the repeated z label runs over
   both trailing axes simultaneously, extracting the diagonal.

The lowering handles static and dynamic (symbolic) shapes, positive/negative
offsets, and arbitrary dim1 / dim2 (including negative indices). Register
"diagonal.default" in ExportedProgramImporter.create_convert_map and
"diagonal" in TorchFXImporter.create_convert_map.

Validation

In-tree regression test (added)

test_einsum_repeated_subscript in
tests/python/relax/test_frontend_from_exported_program.py:

• verify_model against the exact lowering IR for "ii->i" on the default
  decomposition path (this case used to raise the assertion);
• verify_model_numerically for "ii->" (trace), "...ii->...i" (batched
  diagonal), the attention-style two-operand "abca,abcb->c", and the direct
  ops torch.diagonal(x, offset, 0, 1) and torch.trace.

Differential test

verify_patch.py runs on the locked build and simulates the pre-fix behavior
at runtime (popping diagonal.default from the generated convert map):

• Baseline (pre-fix): all 21 diagonal-producing cases (10 issue einsum
  equations + 11 direct torch.diagonal/torch.trace/torch.diag) reproduce
  the exact AssertionError: Unsupported function types ['diagonal.default'];
  1 case (torch.diag on a 1-D input, which goes through diag_embed) is
  unaffected and stays correct in baseline.
• Post-fix: all 22 issue + direct-op cases convert and match PyTorch with
  max|diff| = 0.
• Dynamic shapes: "ii->i" and "...ii->...i" with symbolic dims (both
  diagonal dims sharing one Dim) match PyTorch exactly.
• Regression: the regular einsum family (matmul, transpose, dot, outer,
  batch matmul, ellipsis broadcasting/summation, 3-operand, implicit output) —
  15 cases — all still match with max|diff| = 0.

Run:

TVM_LIBRARY_PATH=<tvm>/build/lib PYTHONPATH=<tvm 源码>/python \
  /home/shenqingchao/miniconda3/envs/tvm23/bin/python \
  results/TVM/deepseek-v4-flash/prove_hum/torch_einsum/verify_patch.py

Files changed

• python/tvm/relax/frontend/torch/base_fx_graph_translator.py — add
  _diagonal (permute_dims + strided_slice crop + einsum ...zz->...z).
• python/tvm/relax/frontend/torch/exported_program_translator.py — register
  "diagonal.default" in the exported-program convert_map.
• python/tvm/relax/frontend/torch/fx_translator.py — register "diagonal" in
  the from_fx convert_map.
• tests/python/relax/test_frontend_from_exported_program.py — add
  test_einsum_repeated_subscript regression coverage.

[tlopex]

Please clamp diag_len to zero. An offset outside the axis range is valid and should produce an empty diagonal. For example, a (3, 4) input with offset=5 returns shape (0,) in PyTorch, but this computes diag_len=-1, causing incompatible slice extents; offset=6 even produces the incorrect shape (1,). Please use tirx.max(0, tirx.min(...)) in both branches and add tests for out-of-range positive and negative offsets.

[siyiweigeHEW]

Thanks for the careful review — you're right, and thanks for the concrete repro.

Confirmed: for a (3, 4) input with offset=5, the lowering computed
diag_len = min(3, 4 - 5) = -1, producing incompatible slice extents (the
...zz->...z einsum then failed with Cannot broadcast extents 2 and 0), and
offset=6 even produced an incorrect non-empty shape.

Fixed in 1eb5cbb by clamping the diagonal length to zero in both branches:

if offset >= 0:
    diag_len = tirx.max(0, tirx.min(n, m - offset))
    ...
else:
    diag_len = tirx.max(0, tirx.min(n + offset, m))

An out-of-range offset now lowers to an empty diagonal of shape (0,),
matching PyTorch, and the negative-branch case (n + offset) is clamped the
same way.

Tests added:

• In-tree: verify_model_numerically in test_einsum_repeated_subscript now
  covers out-of-range offset in {4, 5, 6, -3, -4, -5, -6} on a (3, 4)
  input — each yields (0,) and matches PyTorch exactly.
• Differential suite: SUITE_DIRECT gained 7 out-of-range cases
  ((3, 3) offsets 3, 4, -3; (3, 4) offsets 5, 6, -4, -5).

Verification after the fix: baseline reproduces the original assertion on
28/28 diagonal-producing cases; post-fix 29/29 match PyTorch with
max|diff| = 0; dynamic-shape 3/3; regular einsum regression 15/15.
test_einsum_repeated_subscript passes.