siyiweigeHEW commented on PR #20237: URL: https://github.com/apache/tvm/pull/20237#issuecomment-5499914740
Thanks for the careful review — you're right, and thanks for the concrete repro. Confirmed: for a `(3, 4)` input with `offset=5`, the lowering computed `diag_len = min(3, 4 - 5) = -1`, producing incompatible slice extents (the `...zz->...z` einsum then failed with `Cannot broadcast extents 2 and 0`), and `offset=6` even produced an incorrect non-empty shape. Fixed in [1eb5cbb](https://github.com/apache/tvm/commit/1eb5cbb) by clamping the diagonal length to zero in both branches: ```python if offset >= 0: diag_len = tirx.max(0, tirx.min(n, m - offset)) ... else: diag_len = tirx.max(0, tirx.min(n + offset, m)) ``` An out-of-range offset now lowers to an empty diagonal of shape `(0,)`, matching PyTorch, and the negative-branch case (`n + offset`) is clamped the same way. Tests added: - In-tree: `verify_model_numerically` in `test_einsum_repeated_subscript` now covers out-of-range `offset` in `{4, 5, 6, -3, -4, -5, -6}` on a `(3, 4)` input — each yields `(0,)` and matches PyTorch exactly. - Differential suite: `SUITE_DIRECT` gained 7 out-of-range cases (`(3, 3)` offsets `3, 4, -3`; `(3, 4)` offsets `5, 6, -4, -5`). Verification after the fix: baseline reproduces the original assertion on 28/28 diagonal-producing cases; post-fix 29/29 match PyTorch with `max|diff| = 0`; dynamic-shape 3/3; regular einsum regression 15/15. `test_einsum_repeated_subscript` passes. -- 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]
