(tvm) branch main updated: [Relax][ONNX] Preserve bool dtype when folding constant comparisons (#20286)

[tlopex]

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 2bc2db236d [Relax][ONNX] Preserve bool dtype when folding constant 
comparisons (#20286)
2bc2db236d is described below

commit 2bc2db236d05daedda9674596755a31555cabdb3
Author: T1Gang <143160405+strongbodystrongm...@users.noreply.github.com>
AuthorDate: Tue Sep 8 12:16:17 2026 +0800

    [Relax][ONNX] Preserve bool dtype when folding constant comparisons (#20286)
    
    Fixes #20282.
    
    The ONNX frontend constant-folds binary operators with NumPy. When both
    operands have the same dtype, the result was unconditionally cast back
    to
    the operand dtype to avoid NumPy precision widening.
    
    This is correct for arithmetic operations, but incorrect for comparison
    operators such as `Less`, `LessOrEqual`, `Greater`, and
    `GreaterOrEqual`,
    whose result dtype must be `bool`.
    
    As a result, constant-folded comparisons produced correct 0/1 values but
    returned `int32` or `float32` tensors instead of `bool`.
    
    This patch preserves NumPy boolean results while keeping the existing
    no-widening behavior for non-boolean binary operations.
    
    Tests:
    - Added regression coverage for constant-folded ONNX comparisons with
      `int32` and `float32` operands.
    - The 8 new cases fail before the fix and pass after it.
    - Related ONNX binary tests: 16 passed.
    - `git diff --check upstream/main..HEAD`: passed.
---
 python/tvm/relax/frontend/onnx/onnx_frontend.py |  2 +-
 tests/python/relax/test_frontend_onnx.py        | 50 +++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py 
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 725746c805..2a9cea943c 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -608,7 +608,7 @@ class BinaryBase(OnnxOpConverter):
                 if hasattr(output, "item"):
                     output = output.item()
                 return relax.prim_value(output)
-            if x.dtype == y.dtype:
+            if x.dtype == y.dtype and not _np.issubdtype(output.dtype, 
_np.bool_):
                 # no numpy precision widening
                 output = output.astype(x.dtype)
             if all([isinstance(inp, relax.Constant) for inp in inputs]):
diff --git a/tests/python/relax/test_frontend_onnx.py 
b/tests/python/relax/test_frontend_onnx.py
index 591a2a8ee3..94268d477e 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -582,6 +582,56 @@ def test_concat_with_param_tensor_keeps_runtime_param():
     np.testing.assert_array_equal(params["main"][0].numpy(), weight_np)
 
 
+@pytest.mark.parametrize(
+    "op_name,np_op",
+    [
+        ("Less", np.less),
+        ("LessOrEqual", np.less_equal),
+        ("Greater", np.greater),
+        ("GreaterOrEqual", np.greater_equal),
+    ],
+)
+@pytest.mark.parametrize("np_dtype", ["int32", "float32"])
+def test_constant_comparison_outputs_bool(op_name, np_op, np_dtype):
+    a_np = np.array([[1], [5]], dtype=np_dtype)
+    b_np = np.array([[3]], dtype=np_dtype)
+    rhs_np = np.array([[3]], dtype=np_dtype)
+    graph = helper.make_graph(
+        [
+            helper.make_node("Identity", ["d"], ["dummy"]),
+            helper.make_node("Concat", ["a", "b"], ["lhs"], axis=0),
+            helper.make_node(op_name, ["lhs", "rhs"], ["y"]),
+        ],
+        "constant_comparison",
+        [helper.make_tensor_value_info("d", TensorProto.INT32, [1])],
+        [
+            helper.make_tensor_value_info("y", TensorProto.BOOL, [3, 1]),
+            helper.make_tensor_value_info("dummy", TensorProto.INT32, [1]),
+        ],
+        initializer=[
+            numpy_helper.from_array(a_np, "a"),
+            numpy_helper.from_array(b_np, "b"),
+            numpy_helper.from_array(rhs_np, "rhs"),
+        ],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 
18)], ir_version=9)
+    onnx.checker.check_model(model)
+
+    mod = from_onnx(model, opset=18, shape_dict={"d": [1]}, 
keep_params_in_input=False)
+    constants = []
+
+    def collect_constants(expr):
+        if isinstance(expr, relax.Constant):
+            constants.append(expr.data.numpy())
+
+    relax.analysis.post_order_visit(mod["main"].body, collect_constants)
+    folded_outputs = [arr for arr in constants if arr.shape == (3, 1)]
+    assert len(folded_outputs) == 1
+    expected = np_op(np.concatenate([a_np, b_np], axis=0), rhs_np)
+    np.testing.assert_array_equal(folded_outputs[0], expected)
+    assert folded_outputs[0].dtype == np.dtype("bool")
+
+
 @pytest.mark.parametrize("op_name", ["Add", "Sub", "Mul", "Div", "Pow"])
 def test_binary(op_name: str):
     verify_binary(op_name, [1, 32], [1, 32], [1, 32])