This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e6c6cc8785 [Fix][Relax][ONNX] Support Hardmax over a symbolic axis
extent (#20416)
e6c6cc8785 is described below
commit e6c6cc8785091ad10c18b495a1dbdac6cdb8fc51
Author: Gawr Gura 0721 <[email protected]>
AuthorDate: Fri Sep 25 11:46:01 2026 +0800
[Fix][Relax][ONNX] Support Hardmax over a symbolic axis extent (#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.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 8 ++++++++
tests/python/relax/test_frontend_onnx.py | 13 +++++++++++++
2 files changed, 21 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index a3e6140381..c078073c33 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1080,6 +1080,14 @@ class Hardmax(OnnxOpConverter):
data = bb.normalize(data)
normalized_axis, axis_extent = _get_axis_extent(data, axis, "Hardmax")
dtype = data.ty.dtype
+ if not isinstance(axis_extent, int):
+ # one_hot needs a static depth. For a symbolic extent, mark the
first
+ # maximum by comparing its index against a broadcast iota instead.
+ argmax = relax.op.argmax(data, axis=normalized_axis, keepdims=True)
+ iota_shape = [1] * data.ty.ndim
+ iota_shape[normalized_axis] = axis_extent
+ iota = relax.op.reshape(relax.op.arange(0, axis_extent,
dtype="int64"), iota_shape)
+ return relax.op.astype(relax.op.equal(iota, argmax), dtype)
argmax = relax.op.argmax(data, axis=normalized_axis)
on_value = relax.prim_value(tvm.tirx.const(1.0, dtype))
off_value = relax.prim_value(tvm.tirx.const(0.0, dtype))
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index abef4b7cf5..ec8bc6daed 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1466,6 +1466,19 @@ def test_hardmax_ir():
tvm.ir.assert_structural_equal(tvm_model, Expected)
[email protected]("opset", [11, 13])
[email protected]("axis", [0, 1, -1])
[email protected]("symbolic", [False, True], ids=["static", "symbolic"])
+def test_hardmax(opset, axis, symbolic):
+ shape = [2, 3, 4]
+ model = make_unary_model(
+ "Hardmax", ["N", "C", "W"] if symbolic else shape, attrs={"axis": axis}
+ )
+ # Few distinct values, so rows tie and the first maximum must win.
+ x = rg.integers(0, 3, size=shape).astype("float32")
+ check_correctness(model, inputs={"x": x}, opset=opset)
+
+
def test_legacy_softmax_family_opset11_axis_semantics():
def verify_legacy_softmax_family_axis_ir(op_name: str, expected,
axis_attr: int | None = None):
attrs = {} if axis_attr is None else {"axis": axis_attr}
