LiRunGuo opened a new pull request, #20416: URL: https://github.com/apache/tvm/pull/20416
`Hardmax` is lowered to `argmax` followed by `one_hot`, but `relax.op.one_hot` takes its depth as a static integer attribute. When the reduced axis has a symbolic extent, the importer passed a `tir.Var` as the depth and failed with `TypeError: 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 a `prim.Mul`) whenever any of them is symbolic, for example `Hardmax(axis=1)` on an `[N, 3, W]` input. This keeps the `one_hot` lowering whenever the extent is static, so the IR for existing models is unchanged. For a symbolic extent it compares the `keepdims` argmax index against an `arange` over the axis, broadcast along it, and casts the result to the input dtype. `argmax` returns 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. -- 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]
