tlopex commented on code in PR #20237:
URL: https://github.com/apache/tvm/pull/20237#discussion_r3902530929
##########
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)
Review Comment:
Please clamp `diag_len` to zero. An offset outside the axis range is valid
and should produce an empty diagonal. For example, a `(3, 4)` input with
`offset=5` returns shape `(0,)` in PyTorch, but this computes `diag_len=-1`,
causing incompatible slice extents; `offset=6` even produces the incorrect
shape `(1,)`. Please use `tirx.max(0, tirx.min(...))` in both branches and add
tests for out-of-range positive and negative offsets.
--
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]
