Conversation
Hardmax is lowered to argmax followed by one_hot, but one_hot takes its depth as a static integer attribute. When the reduced axis has a symbolic extent the importer passed a tir.Var (or, for opset <= 12, the symbolic product of the flattened dims) and failed with a TypeError. Keep the one_hot lowering for static extents, and for symbolic extents compare the keepdims argmax against an arange over the axis broadcast along it. argmax returns the first maximum, so ties still resolve to the first index as ONNX requires.
This branch has not been deployed
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.
Hardmaxis lowered toargmaxfollowed byone_hot, butrelax.op.one_hottakes its depth as a static integer attribute. When the reduced axis has a symbolic extent, the importer passed atir.Varas the depth and failed withTypeError: Mismatched type on argument #3 when calling relax.op.one_hot ... Expected int but got ir.Var. For opset 12 and below the input is first flattened to 2-D, so the reduced extent is the product of the trailing dimensions and fails the same way (as aprim.Mul) whenever any of them is symbolic, for exampleHardmax(axis=1)on an[N, 3, W]input.This keeps the
one_hotlowering whenever the extent is static, so the IR for existing models is unchanged. For a symbolic extent it compares thekeepdimsargmax index against anarangeover the axis, broadcast along it, and casts the result to the input dtype.argmaxreturns the first maximum, so ties still resolve to the first index as the ONNX spec requires.The tests cover opsets 11 and 13 with the reduced axis symbolic, a flattened extent that is a product involving a symbolic dim, a static axis on an otherwise symbolic input, and inputs with ties, all compared against onnxruntime. Six of the seven fail without this change; the remaining one checks the unchanged static-extent path.