hiyufan opened a new pull request, #20375: URL: https://github.com/apache/tvm/pull/20375
### What this adds Three reductions that had no converter, so any model using them failed at import with `Unsupported function types`: | op | ATen overload | note | | --- | --- | --- | | `torch.amax(x, dim, keepdim)` | `amax.default` | `torch.logsumexp` decomposes through it, so that was blocked too | | `torch.amin(x, dim, keepdim)` | `amin.default` | | | `torch.min(x, dim, keepdim)` → `(values, indices)` | `min.dim` | `max.dim` was already supported | Found by a differential sweep of the frontend against `torch.export` (88 ops × 7 input shapes, static and dynamic); these accounted for 18 of the 90 programs that raised. ### How - `_amax_amin(op)` → `relax.op.max` / `relax.op.min` over `dim`, a list where empty means every axis (torch's convention; `_sum` handles the same case). Dispatched in both the exported-program and fx translators. - `_max_dim` becomes `_max_min_dim(largest)` and serves both `max.dim` and `min.dim` through `topk(k=1, largest=…)`, the same way `max.dim` was already implemented. No behaviour change for `max.dim`; `test_max_dim` passes unchanged. ### Verification Reduction-op raises in the sweep go from 90 to 72 with this change; the one remaining program per op is the 0-d `()` input, which is a separate issue across all reductions and out of scope here. `test_frontend_from_exported_program.py` and `test_frontend_from_fx.py`: failure sets identical before and after apart from the new tests. `ruff check` / `ruff format --check` (v0.12.3) clean. ### Tests - `test_amax_amin` — IR-level: `amax` over one axis, `amin` over two axes with `keepdim`, `amax` over every axis; numeric checks of the three plus `logsumexp`. - `test_min_dim` — IR-level for `(values, indices)` without keepdim; numeric check of both branches with values *and* indices compared against torch on a permutation input, so the argmin is unambiguous. Both fail against the previous head: ```text test_amax_amin AssertionError: Unsupported function types ['amax.default'] test_min_dim AssertionError: Unsupported function types ['min.dim'] ``` Independent of #20372 / #20373 / #20374 (branched from `main`). --- This change was prepared with AI assistance (Claude). I have reviewed and verified it, and can speak to it in review. -- 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]
