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

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -1264,6 +1264,68 @@ 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).
+        """
+
+        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,

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

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,

tests/python/relax/test_frontend_from_exported_program.py

@@ -3327,6 +3327,79 @@ 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):