This is an automated email from the ASF dual-hosted git repository.
mshr-h 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 0054a95abf [Relax][ONNX] Add Celu operator support (#20923)
0054a95abf is described below
commit 0054a95abf8f0358160f19834a5b8493a3ac4bce
Author: Anirudh Mittal <[email protected]>
AuthorDate: Fri Nov 6 18:03:49 2026 +0530
[Relax][ONNX] Add Celu operator support (#20923)
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`.
Co-authored-by: Anirudh-65 <[email protected]>
---
python/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 8866019c55..2f97ef839a 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1961,6 +1961,22 @@ class Elu(OnnxOpConverter):
) + 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 bba6ef4faf..d34d208b00 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -4480,6 +4480,66 @@ def test_trilu_dynamic_k_ir(upper: bool):
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)
