From 4e8fb28947569f9d70f9c45ef4e64102162c30c0 Mon Sep 17 00:00:00 2001 From: HuEnwei Date: Sun, 30 Aug 2026 18:26:46 +0800 Subject: [PATCH 1/3] [Relax][Frontend][Torch] Support aten.diagonal from decomposed repeated-subscript einsum from_exported_program runs run_decompositions() by default, which lowers torch.einsum with repeated subscripts (diagonal / trace, e.g. "ii->i", "ii->", "...ii->...i") to aten.diagonal + permute (+ sum). The torch frontend had no handler for aten.diagonal.default, so every such valid model failed with `AssertionError: Unsupported function types ['diagonal.default']`. The same root cause blocked torch.diagonal / torch.trace. Add BaseFXGraphImporter._diagonal lowering diagonal(x, offset, dim1, dim2) as permute_dims (move dim1/dim2 to trailing axes) -> two strided_slice (crop each trailing axis to the diagonal length, offset-adjusted) -> relax.op.einsum("...zz->...z"). Handles static and dynamic (symbolic) shapes, positive/negative offsets, and arbitrary dim1/dim2 (incl. negative indices). Register "diagonal.default" in the exported-program convert_map and "diagonal" in the from_fx convert_map. Fixes: #20228 --- .../torch/base_fx_graph_translator.py | 63 +++++++++++++++ .../torch/exported_program_translator.py | 1 + .../tvm/relax/frontend/torch/fx_translator.py | 1 + .../test_frontend_from_exported_program.py | 77 +++++++++++++++++++ 4 files changed, 142 insertions(+) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index d600987cdd7b..e3c565f862e8 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -1264,6 +1264,69 @@ def _einsum(self, node: fx.Node) -> relax.Var: operands = args[1] if isinstance(args[1], torch.Size | tuple | list) else args[1:] return self.block_builder.emit(relax.op.einsum(operands, args[0])) + def _diagonal(self, node: fx.Node) -> relax.Var: + """Convert ``aten.diagonal`` / ``torch.diagonal`` to Relax. + + ``diagonal(input, offset=0, dim1=0, dim2=1)`` extracts the elements + ``input[..., i, i + offset]`` along the ``dim1`` / ``dim2`` axes. It + shows up in the exported graph through ``run_decompositions`` of + ``torch.einsum`` with repeated subscripts (e.g. ``"ii->i"``, + ``"ii->"``, ``"...ii->...i"``), which lower to an ``aten.diagonal`` + followed by a ``sum`` reduction. + + We lower it as: permute ``dim1`` / ``dim2`` to the trailing two axes, + slice each trailing axis to the diagonal length (min of the two + extents, adjusted by ``offset``), and take the diagonal with an + einsum contraction ``...zz->...z`` (the repeated ``z`` label runs + over both trailing axes simultaneously). + """ + import torch # type: ignore + + args = self.retrieve_args(node) + x = args[0] + offset = args[1] if len(args) > 1 else node.kwargs.get("offset", 0) + dim1 = args[2] if len(args) > 2 else node.kwargs.get("dim1", 0) + dim2 = args[3] if len(args) > 3 else node.kwargs.get("dim2", 1) + + shape = self.shape_of(x) + ndim = len(shape.values) + dim1 = dim1 if dim1 >= 0 else ndim + dim1 + dim2 = dim2 if dim2 >= 0 else ndim + dim2 + if dim1 == dim2: + raise ValueError(f"diagonal requires dim1 != dim2, got {dim1} == {dim2}") + + offset = int(offset) + # Move dim1, dim2 to the trailing two axes. + perm = [i for i in range(ndim) if i != dim1 and i != dim2] + [dim1, dim2] + permuted = self.block_builder.emit(relax.op.permute_dims(x, perm)) + + n = shape.values[dim1] + m = shape.values[dim2] + if offset >= 0: + diag_len = tirx.min(n, m - offset) + begin1, end1 = 0, diag_len + begin2, end2 = offset, offset + diag_len + else: + diag_len = tirx.min(n + offset, m) + begin1, end1 = -offset, -offset + diag_len + begin2, end2 = 0, diag_len + + # Crop both diagonal axes to the diagonal length so the einsum ``z`` + # label sees equal extents on both trailing axes. + cropped = self.block_builder.emit( + relax.op.strided_slice( + permuted, axes=[ndim - 2], begin=[begin1], end=[end1], strides=[1] + ) + ) + cropped = self.block_builder.emit( + relax.op.strided_slice( + cropped, axes=[ndim - 1], begin=[begin2], end=[end2], strides=[1] + ) + ) + + # ``...zz -> ...z``: keep every leading axis, contract the diagonal pair. + return self.block_builder.emit(relax.op.einsum([cropped], "...zz->...z")) + def _embedding_impl( self, x, diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index ced0aa7b28bd..e6df019c5c5e 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1894,6 +1894,7 @@ def create_convert_map( "conv3d.default": self._conv3d, "convolution.default": self._convolution, "cross_entropy_loss.default": self._cross_entropy_default, + "diagonal.default": self._diagonal, "einsum.default": self._einsum, "embedding.default": lambda node: self._embedding_impl( self.env[node.args[1]], self.env[node.args[0]] diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index bef0b58f088c..451765011245 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -952,6 +952,7 @@ def create_convert_map( "conv2d": self._conv2d, "conv3d": self._conv3d, "cross_entropy": self._cross_entropy, + "diagonal": self._diagonal, "einsum": self._einsum, "interpolate": self._interpolate, "layer_norm": self._layer_norm, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 7dc3c7356414..15a7e70a8cb5 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -3327,6 +3327,83 @@ def main( verify_model(Einsum2(), example_args, {}, Expected2, run_ep_decomposition=False) +def test_einsum_repeated_subscript(): + """einsum with repeated subscripts (diagonal / trace) on the default + decomposition path. + + ``run_decompositions`` (default) lowers repeated-subscript einsum to + ``aten.diagonal`` + ``permute`` (+ ``sum`` for the trace), which the + frontend converts with the ``_diagonal`` lowering: permute the diagonal + dims to the trailing two axes, slice each to the diagonal length, then an + einsum ``...zz->...z`` extracts the diagonal. This used to raise + ``AssertionError: Unsupported function types ['diagonal.default']``. + """ + + class EinsumDiag(Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.einsum("ii->i", x) + + @tvm.script.ir_module + class Expected: + @R.function + def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple( + R.Tensor((3,), dtype="float32") + ): + with R.dataflow(): + lv: R.Tensor((3, 3), dtype="float32") = R.permute_dims(x, axes=[0, 1]) + lv1: R.Tensor((3, 3), dtype="float32") = R.strided_slice( + lv, (0,), (0,), (3,), (1,), assume_inbound=False + ) + lv2: R.Tensor((3, 3), dtype="float32") = R.strided_slice( + lv1, (1,), (0,), (3,), (1,), assume_inbound=False + ) + lv3: R.Tensor((3,), dtype="float32") = R.einsum((lv2,), subscripts="...zz->...z") + lv4: R.Tensor((3,), dtype="float32") = R.permute_dims(lv3, axes=[0]) + lv5: R.Tensor((3,), dtype="float32") = R.permute_dims(lv4, axes=[0]) + gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv5,) + R.output(gv) + return gv + + example_args = (torch.randn(3, 3, dtype=torch.float32),) + verify_model(EinsumDiag(), example_args, {}, Expected) + + class TraceEinsum(Module): + def forward(self, x): + return torch.einsum("ii->", x) + + class BatchedDiagEinsum(Module): + def forward(self, x): + return torch.einsum("...ii->...i", x) + + class AttentionEinsum(Module): + def forward(self, x, y): + return torch.einsum("abca,abcb->c", x, y) + + verify_model_numerically(TraceEinsum(), (torch.randn(4, 4),)) + verify_model_numerically(BatchedDiagEinsum(), (torch.randn(2, 3, 3),)) + verify_model_numerically( + AttentionEinsum(), (torch.randn(3, 3, 4, 3), torch.randn(3, 3, 4, 3)) + ) + + class DirectDiagonal(Module): + def __init__(self): + super().__init__() + self.offset = 1 + + def forward(self, x): + return torch.diagonal(x, self.offset, 0, 1) + + class DirectTrace(Module): + def forward(self, x): + return torch.trace(x) + + verify_model_numerically(DirectDiagonal(), (torch.randn(3, 4),)) + verify_model_numerically(DirectTrace(), (torch.randn(4, 4),)) + + def test_outer(): class Outer(torch.nn.Module): def forward(self, x, y): From e78a27f73de99796339b11d557f071250f3cfbee Mon Sep 17 00:00:00 2001 From: HuEnwei Date: Sun, 30 Aug 2026 19:02:11 +0800 Subject: [PATCH 2/3] Remove unused import of torch in base_fx_graph_translator Removed unused import statement for torch. --- python/tvm/relax/frontend/torch/base_fx_graph_translator.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index e3c565f862e8..e42af09f56f3 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -1280,7 +1280,6 @@ def _diagonal(self, node: fx.Node) -> relax.Var: einsum contraction ``...zz->...z`` (the repeated ``z`` label runs over both trailing axes simultaneously). """ - import torch # type: ignore args = self.retrieve_args(node) x = args[0] From d519b97b3ce9bc8a89db46c702eb4fbc2ef4e1bc Mon Sep 17 00:00:00 2001 From: HuEnwei Date: Sun, 30 Aug 2026 19:04:10 +0800 Subject: [PATCH 3/3] Refactor main function and streamline einsum calls --- tests/python/relax/test_frontend_from_exported_program.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 15a7e70a8cb5..7341a22187d1 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -3349,9 +3349,7 @@ def forward(self, x): @tvm.script.ir_module class Expected: @R.function - def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple( - R.Tensor((3,), dtype="float32") - ): + def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((3,), dtype="float32")): with R.dataflow(): lv: R.Tensor((3, 3), dtype="float32") = R.permute_dims(x, axes=[0, 1]) lv1: R.Tensor((3, 3), dtype="float32") = R.strided_slice( @@ -3384,9 +3382,7 @@ def forward(self, x, y): verify_model_numerically(TraceEinsum(), (torch.randn(4, 4),)) verify_model_numerically(BatchedDiagEinsum(), (torch.randn(2, 3, 3),)) - verify_model_numerically( - AttentionEinsum(), (torch.randn(3, 3, 4, 3), torch.randn(3, 3, 4, 3)) - ) + verify_model_numerically(AttentionEinsum(), (torch.randn(3, 3, 4, 3), torch.randn(3, 3, 4, 3))) class DirectDiagonal(Module): def __init__(self):