From 75000ad6c6cf6def9b153481982f9d193ba02cc4 Mon Sep 17 00:00:00 2001 From: tsu-bin Date: Mon, 3 Jun 2024 16:54:49 +0800 Subject: [PATCH] [Relax] [ONNX] Add support for Sign and Not --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 18 ++++++++++++++++++ tests/python/relax/test_frontend_onnx.py | 8 ++++++++ 2 files changed, 26 insertions(+) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 3a70cd090a54..85d4402d6640 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1948,6 +1948,22 @@ def _impl_v14(cls, bb, inputs, attr, params): ) +class Sign(OnnxOpConverter): + """Converts an onnx Sign node into an equivalent Relax expression.""" + + @classmethod + def _impl_v9(cls, bb, inputs, attr, params): + return relax.op.sign(inputs[0]) + + +class Not(OnnxOpConverter): + """Converts an onnx Not node into an equivalent Relax expression.""" + + @classmethod + def _impl_v1(cls, bb, inputs, attr, params): + return relax.op.logical_not(inputs[0]) + + def _get_convert_map(): return { "MatMul": MatMul, @@ -2030,6 +2046,8 @@ def _get_convert_map(): "Elu": Elu, "HardSigmoid": HardSigmoid, "HardSwish": HardSwish, + "Sign": Sign, + "Not": Not, } diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 0fc7ec064402..05316f2699dd 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -600,6 +600,14 @@ def test_hardswish(): verify_unary("HardSwish", [32, 32]) +def test_sign(): + verify_unary("Sign", [32, 32]) + + +def test_not(): + verify_unary("Not", [32, 32], dtype=TensorProto.BOOL) + + def test_conv(): def _verify_conv(input_shape, weight_shape, output_shape): bias_shape = [output_shape[1]]