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 513756327d [Relax][ONNX] Support lower-rank PRelu slopes (#20115)
513756327d is described below
commit 513756327d14bdd53d42d1e591a797eda723cda7
Author: Hongyi Wu <[email protected]>
AuthorDate: Wed Aug 12 02:45:13 2026 +0800
[Relax][ONNX] Support lower-rank PRelu slopes (#20115)
## Summary
This PR extends the Relax ONNX `PRelu` converter to support lower-rank
slope tensors that ONNX aligns to the trailing dimensions of the input.
The gap was exposed by Qualcomm's Real-ESRGAN-General-x4v3 export. Its
activation input has shape `[1, 64, 128, 128]`, while its slope has
shape `[64, 1, 1]`. This is valid ONNX unidirectional broadcasting, but
the current converter rejects it because the two ranks differ.
## Goal
Import legal lower-rank ONNX `PRelu` slopes when they can be represented
by Relax's one-dimensional `nn.prelu` slope and an adjusted axis.
## What changed
- Align lower-rank slopes to the trailing input dimensions.
- Translate the slope's non-broadcast dimension to the corresponding
Relax input axis.
- Keep rejecting slopes with multiple non-broadcast dimensions, which
cannot be represented by the current Relax `nn.prelu` operator.
- Handle rank-zero slopes without indexing an empty shape.
- Add structural and ONNX Runtime-backed numerical regression coverage.
## Design
For a slope with at most one non-broadcast dimension, let
`relative_axis` be that dimension in the slope. ONNX trailing-dimension
alignment maps it to:
```text
axis = input_rank - slope_rank + relative_axis
```
For the motivating shape pair:
```text
input: [1, 64, 128, 128]
slope: [64, 1, 1]
aligned slope: [1, 64, 1, 1]
Relax axis: 1
```
The converter then reshapes the slope to `[64]` and emits
`R.nn.prelu(..., axis=1)`.
## Updated converter behavior
| ONNX slope shape | Behavior |
| --- | --- |
| Rank-zero or all-one shape | Reshape to a one-element vector |
| Rank-one shape | Preserve the existing final-axis behavior |
| Lower/equal rank with one non-broadcast dimension | Align to trailing
input dimensions and emit the corresponding Relax axis |
| Multiple non-broadcast dimensions | Continue to raise an explicit
unsupported-shape error |
## Safety checks
- Existing scalar, one-dimensional, and same-rank structural cases
remain covered.
- The new structural case checks input `[1, 32, 16, 16]`, slope `[32, 1,
1]`, and Relax `axis=1`.
- The new numerical case compares TVM with ONNX Runtime using
channel-specific negative slopes.
- The full Relax ONNX frontend test file passes in the validation
environment, apart from five pre-existing Float8 baseline cases that
were excluded explicitly.
## Out of scope / non-goals
- Supporting arbitrary slopes with multiple non-broadcast dimensions.
- Changing Relax `nn.prelu` semantics or legalization.
- Adding the external Real-ESRGAN model to the TVM test suite.
## Results
The pinned Real-ESRGAN-General-x4v3 ONNX model contains 33 `PRelu`
nodes. With this change it imports, compiles for the C target, and runs
end to end:
```text
input: [1, 3, 128, 128]
output: [1, 3, 512, 512]
max abs error: 4.0531158447265625e-06 versus ONNX Runtime
```
## Tests
- `pre-commit run --files
python/tvm/relax/frontend/onnx/onnx_frontend.py
tests/python/relax/test_frontend_onnx.py`
- Focused H20 run: `2 passed, 498 deselected`
- Relax ONNX frontend H20 run: `482 passed, 9 skipped, 5 deselected, 4
xfailed`
- Pinned Real-ESRGAN-General-x4v3 end-to-end C-target validation against
ONNX Runtime
## References
- [ONNX PRelu
specification](https://onnx.ai/onnx/operators/onnx__PRelu.html)
- [Qualcomm
Real-ESRGAN-General-x4v3](https://huggingface.co/qualcomm/Real-ESRGAN-General-x4v3/tree/e12a7dcde3df0cf4315c648e0b5e4ca4f43d6904)
- [Previous Relax ONNX PRelu
support](https://github.com/apache/tvm/pull/18658)
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 13 ++++--
tests/python/relax/test_frontend_onnx.py | 60 +++++++++++++++++++++++--
2 files changed, 66 insertions(+), 7 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index c7e6b5b585..6bbf220dfe 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1777,11 +1777,15 @@ class PRelu(OnnxOpConverter):
ndim = len(x_shape)
s_ndim = len(slope_shape)
- if all(ss == 1 for ss in slope_shape) or s_ndim == 1:
+ if all(ss == 1 for ss in slope_shape):
+ slope = relax.op.reshape(slope, (1,))
+ return relax.op.nn.prelu(x, slope, ndim - 1)
+
+ if s_ndim == 1:
slope = relax.op.reshape(slope, (slope_shape[0],))
return relax.op.nn.prelu(x, slope, ndim - 1)
- if s_ndim == ndim:
+ 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
@@ -1789,9 +1793,10 @@ class PRelu(OnnxOpConverter):
raise ValueError(
f"Invalid PRelu slope shape (multiple non-broadcast dims):
{slope_shape}"
)
- axis = non_one_axes[0]
+ relative_axis = non_one_axes[0]
+ axis = ndim - s_ndim + relative_axis
- slope = relax.op.reshape(slope, (slope_shape[axis],))
+ slope = relax.op.reshape(slope, (slope_shape[relative_axis],))
return relax.op.nn.prelu(x, slope, axis)
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 439c4374c5..4adc3a0ab3 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3455,21 +3455,35 @@ def test_mish():
def test_prelu():
- def _assert_prelu_ir(slope_shape, expected):
+ def _assert_prelu_ir(slope_shape, expected, input_shape=(3, 32, 32)):
prelu_node = helper.make_node("PRelu", ["a", "b"], ["c"])
graph = helper.make_graph(
[prelu_node],
"prelu_structural_test",
inputs=[
- helper.make_tensor_value_info("a", TensorProto.FLOAT, [3, 32,
32]),
+ helper.make_tensor_value_info("a", TensorProto.FLOAT,
input_shape),
helper.make_tensor_value_info("b", TensorProto.FLOAT,
slope_shape),
],
- outputs=[helper.make_tensor_value_info("c", TensorProto.FLOAT, [3,
32, 32])],
+ outputs=[helper.make_tensor_value_info("c", TensorProto.FLOAT,
input_shape)],
)
model = helper.make_model(graph, producer_name="prelu_structural_test")
tvm_model = from_onnx(model, keep_params_in_input=True)
tvm.ir.assert_structural_equal(tvm_model, expected)
+ @I.ir_module
+ class ExpectedRankZeroSlope:
+ @R.function
+ def main(
+ a: R.Tensor((3, 32, 32), dtype="float32"),
+ b: R.Tensor((), dtype="float32"),
+ ) -> R.Tensor((3, 32, 32), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ lv: R.Tensor((1,), dtype="float32") = R.reshape(b,
R.shape([1]))
+ gv: R.Tensor((3, 32, 32), dtype="float32") = R.nn.prelu(a, lv,
axis=2)
+ R.output(gv)
+ return gv
+
@I.ir_module
class ExpectedScalarSlope:
@R.function
@@ -3526,10 +3540,50 @@ def test_prelu():
R.output(gv)
return gv
+ @I.ir_module
+ class ExpectedLowerRankChannelSlope:
+ @R.function
+ def main(
+ a: R.Tensor((1, 32, 16, 16), dtype="float32"),
+ b: R.Tensor((32, 1, 1), dtype="float32"),
+ ) -> R.Tensor((1, 32, 16, 16), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ lv: R.Tensor((32,), dtype="float32") = R.reshape(b,
R.shape([32]))
+ gv: R.Tensor((1, 32, 16, 16), dtype="float32") = R.nn.prelu(a,
lv, axis=1)
+ 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))
+
+
+def test_prelu_lower_rank_slope():
+ input_shape = (1, 4, 3, 3)
+ slope_shape = (4, 1, 1)
+ graph = helper.make_graph(
+ [helper.make_node("PRelu", ["x", "slope"], ["y"])],
+ "prelu_lower_rank_slope_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
+ helper.make_tensor_value_info("slope", TensorProto.FLOAT,
slope_shape),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
input_shape)],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="prelu_lower_rank_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),
+ "slope": np.array([0.1, 0.2, 0.3, 0.4],
dtype="float32").reshape(slope_shape),
+ }
+ check_correctness(model, inputs=inputs, opset=16, check_dtypes=True)
def test_thresholded_relu():
