siyiweigeHEW opened a new pull request, #20241:
URL: https://github.com/apache/tvm/pull/20241
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:
```bash
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.
--
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]
