This is an automated email from the ASF dual-hosted git repository.
tlopex 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 38fd68c3c2 [Fix][Relax][Frontend][ONNX] Support broadcastable
multi-axis PRelu slopes (#20149)
38fd68c3c2 is described below
commit 38fd68c3c25294c6f4e58dd4b189559abb51da02
Author: HuEnwei <[email protected]>
AuthorDate: Wed Aug 26 06:51:32 2026 +0800
[Fix][Relax][Frontend][ONNX] Support broadcastable multi-axis PRelu slopes
(#20149)
Fixes: #20148
## Summary
The Relax ONNX frontend `PRelu._impl_v1` rejected a valid `PRelu` whose
`slope`
is broadcastable to `X` across **multiple** non-broadcast axes (e.g. a
slope
shaped exactly like `X`, or with several non-1 dims) with a
`ValueError`, even
though onnxruntime and `onnx.reference` accept and run such models
correctly.
After #20115 added lower-rank and rank-0 slope support, this was the
remaining
gap for unidirectionally broadcastable slopes.
## Root cause
`PRelu._impl_v1` lowers a slope with a **single** non-broadcast axis to
`relax.op.nn.prelu(x, slope_vec, axis)`, which is the only shape Relax's
`nn.prelu` can express (one per-axis slope vector). A slope that
broadcasts
along several axes (`slope == X`, `(1, C, H, W)`, `(N, C, 1, 1)`, …)
cannot be
represented that way, so the code raised instead of lowering the op:
```python
# Must have only ONE non-broadcast axis
if len(non_one_axes) != 1:
raise ValueError(
f"Invalid PRelu slope shape (multiple non-broadcast dims):
{slope_shape}"
)
```
## Fix
For the multi-non-broadcast case, lower `PRelu` elementwise instead:
```python
dtype = x.ty.dtype.dtype
return relax.op.where(
relax.op.less(x, relax.const(0, dtype)),
relax.op.multiply(x, slope),
x,
)
```
This is exactly `PRelu(x, s) = where(x < 0, s * x, x)` with `s`
broadcast to
`x`. The all-ones / rank-1 / single-non-broadcast paths are unchanged
and still
use `nn.prelu`. Non-broadcastable slopes (rank > x, or a non-1 dim that
does not
match `x`) still fail at build time with a broadcast error, consistent
with
onnxruntime rejecting them at session creation.
## Validation
Differential test: Relax (build + `VirtualMachine`) vs onnxruntime on
the same
model, comparing output shapes and max |diff| (cross-checked with
`onnx.reference`), across 7 `X` shapes × every valid broadcastable
`slope`.
| Case | onnxruntime | TVM | max\|diff\| | Result |
|---|---|---|---|---|
| lower-rank `(64,1,1)` on `(1,64,128,128)` | `(1,64,128,128)` |
`(1,64,128,128)` | 0.00e+00 | OK |
| rank-0 scalar slope | `(2,3,4,5)` | `(2,3,4,5)` | 0.00e+00 | OK |
| same-rank multi non-broadcast, `slope == X` | `(2,3,4,5)` |
`(2,3,4,5)` | 0.00e+00 | OK |
| same-rank multi non-broadcast `(1,3,4,5)` | `(2,3,4,5)` | `(2,3,4,5)`
| 0.00e+00 | OK |
| lower-rank multi non-broadcast `(4,5)` | `(2,3,4,5)` | `(2,3,4,5)` |
0.00e+00 | OK |
| same-rank single non-broadcast `(1,1,4,1)` (regression) | `(2,3,4,5)`
| `(2,3,4,5)` | 0.00e+00 | OK |
| 1-D `(5,)` / all-ones (regression) | `(2,3,4,5)` | `(2,3,4,5)` |
0.00e+00 | OK |
Total: **113 valid cases (onnxruntime-accepted), 113 OK, 0 rejected, 0
numeric
mismatches** (pre-fix: 46 OK / 67 rejected). The motivating Real-ESRGAN
case
`X(1,64,128,128)` + `slope(64,1,1)` now imports and matches onnxruntime
exactly.
Run:
```bash
python results/TVM/deepseek-v4-flash/prove_hum/onnx_PRelu/5修复_差分验证.py
```
## Files changed
- `python/tvm/relax/frontend/onnx/onnx_frontend.py` — `PRelu._impl_v1`:
elementwise
`where(x < 0, s * x, x)` fallback for multi-non-broadcast slopes.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 28 +++++++++------
tests/python/relax/test_frontend_onnx.py | 46 +++++++++++++++++++++++++
2 files changed, 64 insertions(+), 10 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index b63518f6c0..054efa9f0c 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1802,16 +1802,24 @@ class PRelu(OnnxOpConverter):
if s_ndim <= ndim:
non_one_axes = [i for i, ss in enumerate(slope_shape) if ss != 1]
- # Must have only ONE non-broadcast axis
- if len(non_one_axes) != 1:
- raise ValueError(
- f"Invalid PRelu slope shape (multiple non-broadcast dims):
{slope_shape}"
- )
- relative_axis = non_one_axes[0]
- axis = ndim - s_ndim + relative_axis
-
- slope = relax.op.reshape(slope, (slope_shape[relative_axis],))
- return relax.op.nn.prelu(x, slope, axis)
+ # A single non-broadcast axis can be expressed directly as a
+ # per-axis slope of nn.prelu.
+ if len(non_one_axes) == 1:
+ relative_axis = non_one_axes[0]
+ axis = ndim - s_ndim + relative_axis
+
+ slope = relax.op.reshape(slope, (slope_shape[relative_axis],))
+ return relax.op.nn.prelu(x, slope, axis)
+
+ # Multiple non-broadcast axes (including a slope shaped like x):
+ # nn.prelu can only express a single per-axis slope, so lower
+ # PRelu(x, s) = where(x < 0, s * x, x) elementwise instead.
+ dtype = x.ty.dtype.dtype
+ return relax.op.where(
+ relax.op.less(x, relax.const(0, dtype)),
+ relax.op.multiply(x, slope),
+ x,
+ )
raise ValueError(f"Unsupported PRelu slope shape: {slope_shape}")
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 39e067d5d7..60046d44df 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3727,12 +3727,28 @@ def test_prelu():
R.output(gv)
return gv
+ @I.ir_module
+ class ExpectedMultiAxisSlope:
+ @R.function
+ def main(
+ a: R.Tensor((2, 3, 4, 5), dtype="float32"),
+ b: R.Tensor((4, 5), dtype="float32"),
+ ) -> R.Tensor((2, 3, 4, 5), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ lv: R.Tensor((2, 3, 4, 5), dtype="bool") = R.less(a,
R.const(0.0, "float32"))
+ lv1: R.Tensor((2, 3, 4, 5), dtype="float32") = R.multiply(a, b)
+ gv: R.Tensor((2, 3, 4, 5), dtype="float32") = R.where(lv, lv1,
a)
+ R.output(gv)
+ return gv
+
_assert_prelu_ir([], ExpectedRankZeroSlope)
_assert_prelu_ir([1], ExpectedScalarSlope)
_assert_prelu_ir([1, 1], ExpectedTwoDimScalarSlope)
_assert_prelu_ir([32], ExpectedChannelSlope)
_assert_prelu_ir([3, 1, 1], ExpectedBatchSlope)
_assert_prelu_ir([32, 1, 1], ExpectedLowerRankChannelSlope,
input_shape=(1, 32, 16, 16))
+ _assert_prelu_ir([4, 5], ExpectedMultiAxisSlope, input_shape=(2, 3, 4, 5))
def test_prelu_lower_rank_slope():
@@ -3759,6 +3775,36 @@ def test_prelu_lower_rank_slope():
check_correctness(model, inputs=inputs, opset=16, check_dtypes=True)
+def test_prelu_multi_axis_slope():
+ """A slope broadcastable across multiple axes (incl. a slope shaped like
x) is
+ lowered elementwise to PRelu(x, s) = where(x < 0, s * x, x) since nn.prelu
can
+ only express a single per-axis slope."""
+ input_shape = (2, 3, 4, 5)
+ for slope_shape in [(4, 5), (1, 3, 4, 5), (2, 3, 4, 5)]:
+ graph = helper.make_graph(
+ [helper.make_node("PRelu", ["x", "slope"], ["y"])],
+ "prelu_multi_axis_slope_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT,
input_shape),
+ helper.make_tensor_value_info("slope", TensorProto.FLOAT,
list(slope_shape)),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
input_shape)],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="prelu_multi_axis_slope_test",
+ opset_imports=[helper.make_opsetid("", 16)],
+ )
+ inputs = {
+ "x": np.linspace(-2.0, 2.0, np.prod(input_shape),
dtype="float32").reshape(input_shape),
+ # negative slopes exercise the s * x path on both sides of the
sign.
+ "slope": np.linspace(-0.5, 0.8, np.prod(slope_shape),
dtype="float32").reshape(
+ slope_shape
+ ),
+ }
+ check_correctness(model, inputs=inputs, opset=16, check_dtypes=True)
+
+
def test_thresholded_relu():
model = make_unary_model("ThresholdedRelu", [2, 3])
tvm_model = from_onnx(model, keep_params_in_input=True)
