[Relax][ONNX] Preserve bool dtype when folding constant comparisons

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -608,7 +608,7 @@ def base_impl(cls, bb, inputs, attr, params):
                 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]):

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])