[Relax][Frontend][Torch] Clamp diagonal length to zero for out-of-range offsets

For an out-of-range offset (|offset| >= max(extent1, extent2)), PyTorch's
torch.diagonal returns an empty diagonal of shape (0,). The lowering
computed diag_len = min(extent1, extent2 - offset), which could go negative:
e.g. a (3, 4) input with offset=5 gave diag_len=-1 and incompatible slice
extents (the einsum then failed to broadcast extents 2 and 0), and offset=6
even produced an incorrect non-empty shape.

Clamp diag_len with tirx.max(0, ...) in both the positive- and
negative-offset branches so an out-of-range offset lowers to an empty
diagonal, matching PyTorch. Add in-tree regression coverage for out-of-range
positive and negative offsets.

Author: siyiweigeHEW

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

@@ -1302,11 +1302,11 @@ def _diagonal(self, node: fx.Node) -> relax.Var:
         n = shape.values[dim1]
         m = shape.values[dim2]
         if offset >= 0:
-            diag_len = tirx.min(n, m - offset)
+            diag_len = tirx.max(0, tirx.min(n, m - offset))
             begin1, end1 = 0, diag_len
             begin2, end2 = offset, offset + diag_len
         else:
-            diag_len = tirx.min(n + offset, m)
+            diag_len = tirx.max(0, tirx.min(n + offset, m))
             begin1, end1 = -offset, -offset + diag_len
             begin2, end2 = 0, diag_len
 

tests/python/relax/test_frontend_from_exported_program.py

@@ -3399,6 +3399,21 @@ def forward(self, x):
     verify_model_numerically(DirectDiagonal(), (torch.randn(3, 4),))
     verify_model_numerically(DirectTrace(), (torch.randn(4, 4),))
 
+    # Out-of-range offsets (|offset| >= max(extent1, extent2)) are valid in
+    # PyTorch and yield an empty diagonal of shape (0,); the lowering must
+    # clamp the diagonal length to zero instead of producing negative slice
+    # extents or a wrong non-empty shape.
+    class DirectDiagonalOutOfRange(Module):
+        def __init__(self, offset):
+            super().__init__()
+            self.offset = offset
+
+        def forward(self, x):
+            return torch.diagonal(x, self.offset, 0, 1)
+
+    for offset in [4, 5, 6, -3, -4, -5, -6]:
+        verify_model_numerically(DirectDiagonalOutOfRange(offset), (torch.randn(3, 4),))
+
 
 def test_outer():
     class Outer(torch.nn.Module):