From c2bdb43d7efd4ca8c5baf4458fe31e0a0b9a5ce4 Mon Sep 17 00:00:00 2001 From: Anirudh-65 Date: Fri, 6 Nov 2026 00:40:38 +0530 Subject: [PATCH] [Relax][ONNX] Add Celu operator support This commit adds support for the ONNX Celu (Continuously Differentiable Exponential Linear Unit) operator in the Relax ONNX frontend. Celu is computed element-wise as: celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1)) = -alpha * relu(1 - exp(x / alpha)) + relu(x) Changes include: - Added `Celu` converter class in `python/tvm/relax/frontend/onnx/onnx_frontend.py` with `_impl_v12` supporting the `alpha` attribute (defaulting to 1.0). - Registered "Celu" in the operator converter map. - Added structural IR equivalence tests (`test_celu_ir`) and numerical correctness tests against ONNX Runtime (`test_celu`) in `tests/python/relax/test_frontend_onnx.py`. --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 17 ++++++ tests/python/relax/test_frontend_onnx.py | 60 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 8866019c5523..2f97ef839a00 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1961,6 +1961,22 @@ def _impl_v1(cls, bb, inputs, attr, params): ) + relax.op.nn.relu(inputs[0]) +class Celu(OnnxOpConverter): + """Converts an onnx Celu node into an equivalent Relax expression. + + celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1)) + = relu(x) - alpha * relu(1 - exp(x / alpha)) + """ + + @classmethod + def _impl_v12(cls, bb, inputs, attr, params): + alpha = float(attr.get("alpha", 1.0)) + x = inputs[0] + return relax.expr.const(-alpha) * relax.op.nn.relu( + relax.expr.const(1.0) - relax.op.exp(x / relax.expr.const(alpha)) + ) + relax.op.nn.relu(x) + + class Selu(OnnxOpConverter): """Converts an onnx Selu node into an equivalent Relax expression.""" @@ -6164,6 +6180,7 @@ def _get_convert_map(): "LeakyRelu": LeakyRelu, "ThresholdedRelu": ThresholdedRelu, "Elu": Elu, + "Celu": Celu, "Gelu": Gelu, "FastGelu": FastGelu, "BiasGelu": BiasGelu, diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index bba6ef4faf0e..d34d208b001c 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -4480,6 +4480,66 @@ def main(x: R.Tensor((m, 3), dtype="float32")) -> R.Tensor((m, 3), dtype="float3 tvm_tirx.ir.assert_structural_equal(tvm_model, expected) +def test_celu_ir(): + """Check the lowered Relax IR for the Celu operator (opset 12, default alpha).""" + model = make_unary_model("Celu", [2, 3]) + model.opset_import[0].version = 12 + tvm_model = from_onnx(model, opset=12, keep_params_in_input=True) + + @I.ir_module + class Expected: + @R.function + def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="float32") = R.divide(x, R.const(1.0, "float32")) + lv1: R.Tensor((2, 3), dtype="float32") = R.exp(lv) + lv2: R.Tensor((2, 3), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv1) + lv3: R.Tensor((2, 3), dtype="float32") = R.nn.relu(lv2) + lv4: R.Tensor((2, 3), dtype="float32") = R.multiply(R.const(-1.0, "float32"), lv3) + lv5: R.Tensor((2, 3), dtype="float32") = R.nn.relu(x) + gv: R.Tensor((2, 3), dtype="float32") = R.add(lv4, lv5) + R.output(gv) + return gv + + tvm_tirx.ir.assert_structural_equal(tvm_model, Expected) + + # Custom alpha=0.5 + model = make_unary_model("Celu", [2, 3], attrs={"alpha": 0.5}) + model.opset_import[0].version = 12 + tvm_model = from_onnx(model, opset=12, keep_params_in_input=True) + + @I.ir_module + class ExpectedCustomAlpha: + @R.function + def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"): + R.func_attr({"num_input": 1}) + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="float32") = R.divide(x, R.const(0.5, "float32")) + lv1: R.Tensor((2, 3), dtype="float32") = R.exp(lv) + lv2: R.Tensor((2, 3), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv1) + lv3: R.Tensor((2, 3), dtype="float32") = R.nn.relu(lv2) + lv4: R.Tensor((2, 3), dtype="float32") = R.multiply(R.const(-0.5, "float32"), lv3) + lv5: R.Tensor((2, 3), dtype="float32") = R.nn.relu(x) + gv: R.Tensor((2, 3), dtype="float32") = R.add(lv4, lv5) + R.output(gv) + return gv + + tvm_tirx.ir.assert_structural_equal(tvm_model, ExpectedCustomAlpha) + + +def test_celu(): + """Check numerical correctness of the Celu operator against onnxruntime.""" + # Default alpha=1.0 + check_correctness(make_unary_model("Celu", [2, 3]), opset=12) + + # Custom alpha=0.5 + check_correctness(make_unary_model("Celu", [2, 3], attrs={"alpha": 0.5}), opset=12) + + # Custom alpha=2.0 + check_correctness(make_unary_model("Celu", [2, 3], attrs={"alpha": 2.0}), opset=12) + + def test_selu(): model = make_unary_model("Selu", [2, 3]) tvm_model = from_onnx(model, keep_params_in_input=True)