Conversation
`torch.amax` / `torch.amin` (`amax.default`, `amin.default`) and the `(values, indices)` form of `torch.min(x, dim)` (`min.dim`) had no converter, so a model using any of them failed at import with `Unsupported function types`. `max.dim` was already supported, and `logsumexp` decomposes through `amax`, so it was blocked by the same gap. `_amax_amin` maps to `relax.op.max` / `relax.op.min` over `dim` (a list; empty means every axis, as in torch). `_max_dim` becomes `_max_min_dim(largest)` and serves both `max.dim` and `min.dim` through `topk(k=1)`, exactly as `max.dim` was already handled. `amax` / `amin` are also dispatched in the fx translator. In the differential sweep of the frontend against torch.export (88 ops x 7 input shapes x static/dynamic), the reduction ops that raised drop from 90 to 72 programs; the ones left for these ops are the 0-d input case, which is a separate, op-wide issue. Tests: IR-level checks for amax over one axis, amin over two axes with keepdim, and amax over every axis; numeric checks of those three plus logsumexp; and min.dim with and without keepdim, values and indices compared with torch on a permutation input so the argmin is unambiguous. Both new tests fail against the previous head with `Unsupported function types ['amax.default']` / `['min.dim']`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this adds
Three reductions that had no converter, so any model using them failed at import with
Unsupported function types:torch.amax(x, dim, keepdim)amax.defaulttorch.logsumexpdecomposes through it, so that was blocked tootorch.amin(x, dim, keepdim)amin.defaulttorch.min(x, dim, keepdim)→(values, indices)min.dimmax.dimwas already supportedFound 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.minoverdim, a list where empty means every axis (torch's convention;_sumhandles the same case). Dispatched in both the exported-program and fx translators._max_dimbecomes_max_min_dim(largest)and serves bothmax.dimandmin.dimthroughtopk(k=1, largest=…), the same waymax.dimwas already implemented. No behaviour change formax.dim;test_max_dimpasses 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.pyandtest_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:amaxover one axis,aminover two axes withkeepdim,amaxover every axis; numeric checks of the three pluslogsumexp.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:
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.