[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.