From 236ab026ed3f6922a34fa1dcb57d25bfdbd7a9c4 Mon Sep 17 00:00:00 2001 From: FFChopon <3483776996@qq.com> Date: Tue, 18 Aug 2026 04:04:42 +0800 Subject: [PATCH] [Relax][Frontend][ONNX] Support broadcastable multi-axis PRelu slopes PRelu._impl_v1 only accepted a slope with a single non-broadcast axis. A valid ONNX PRelu whose slope is broadcastable to x across multiple axes (e.g. a slope shaped like x, or with several non-broadcast dims) was rejected with ValueError 'Invalid PRelu slope shape (multiple non-broadcast dims)'. onnxruntime and onnx.reference both accept and run such models correctly. nn.prelu can only express a single per-axis slope, so lower the multi-axis case elementwise as PRelu(x, s) = where(x < 0, s * x, x) via relax.op.where/less/multiply. The single non-broadcast axis and all-ones/slope-of-rank-1 cases keep using nn.prelu unchanged. Together with #20115 (lower-rank and rank-0 slopes), all unidirectionally broadcastable slope shapes accepted by onnxruntime now import. Co-Authored-By: Claude --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 6bbf220dfe2e..b10f3ab5db4e 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1788,16 +1788,24 @@ def _impl_v1(cls, bb, inputs, attr, params): 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}")