diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 725746c80542..2a9cea943cf7 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/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]): diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 591a2a8ee349..94268d477ee2 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])