[Bug][Relax][Frontend][ONNX] PRelu with lower-rank broadcastable `slope` rejected by `from_onnx`: ValueError "Unsupported PRelu slope shape"

[siyiweigeHEW]

Expected behavior

A valid ONNX PRelu model whose slope is lower-rank than X but unidirectionally broadcastable (numpy trailing-dimension alignment) should be imported successfully by tvm.relax.frontend.onnx.from_onnx. The ONNX spec explicitly permits this — "The shape of slope can be smaller than first input X; if so, its shape must be unidirectional broadcastable to X" (see ONNX PRelu spec). Such models pass onnx.checker and run correctly in onnxruntime and onnx.reference.

Actual behavior

from_onnx raises ValueError: Unsupported PRelu slope shape:

ValueError: Unsupported PRelu slope shape: R.shape([64, 1, 1])

raised at python/tvm/relax/frontend/onnx/onnx_frontend.py:1154 in PRelu._impl_v1 (class at line 1120). The frontend only supports:

• all-ones slope, or rank-1 slope (onnx_frontend.py:1137-1139, reshaped to 1-D and applied at axis = ndim - 1), and
• a same-rank slope with exactly one non-broadcast axis (onnx_frontend.py:1141-1151).

Any lower-rank broadcastable slope (e.g. (64, 1, 1) for X: (1, 64, 128, 128)) falls through to the unconditional raise ValueError at line 1154. Two related coverage gaps in the same method:

• same-rank slope with multiple non-broadcast dims (including slope shaped identically to X) → ValueError: Invalid PRelu slope shape (multiple non-broadcast dims) at line 1147;
• scalar (rank-0) slope → crash IndexError: ShapeExpr index out of range, because slope_shape[0] at line 1138 indexes an empty shape.

The same model passes onnx.checker, runs correctly in onnxruntime (output shape (1, 64, 128, 128)), and is confirmed valid by the reference implementation onnx.reference, so the rejection is a frontend coverage gap, not an invalid model.

Environment

• OS: Linux
• TVM: v0.24.dev0 (main branch, commit 262c6d2e0, built 2026-02-11)
• Python: 3.11
• onnx: 1.20.1
• onnxruntime: 1.24.1

Steps to reproduce

"""Repro: valid ONNX PRelu with lower-rank broadcastable slope is rejected by the
TVM relax ONNX frontend, while onnxruntime accepts and runs it correctly.
This is exactly the motivating case (X(1,64,128,128) + slope(64,1,1)) reported in
apache/tvm #20115, triggered by a Qualcomm Real-ESRGAN export."""
import numpy as np
import onnx, onnxruntime
from onnx import helper, TensorProto
from tvm.relax.frontend.onnx import from_onnx

x_shape, slope_shape = (1, 64, 128, 128), (64, 1, 1)          # slope is lower-rank, broadcastable
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, list(x_shape))
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, list(x_shape))
node = helper.make_node("PRelu", ["X", "slope"], ["Y"])
slope = helper.make_tensor("slope", TensorProto.FLOAT, list(slope_shape),
                           np.random.RandomState(0).randn(*slope_shape).astype("float32").flatten().tolist())
graph = helper.make_graph([node], "prelu", [X], [Y], initializer=[slope])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
model.ir_version = 8

onnx.checker.check_model(model)                                # (1) valid ONNX model
x = np.random.RandomState(1).randn(*x_shape).astype("float32")
y = onnxruntime.InferenceSession(model.SerializeToString()).run(None, {"X": x})[0]
print("onnxruntime runs OK ->", y.shape)                       # (2) reference impl works
from_onnx(model, shape_dict={"X": x_shape})                    # (3) TVM rejects the same model

Actual output:

onnxruntime runs OK -> (1, 64, 128, 128)
Traceback (most recent call last):
  ...
  File "tvm/relax/frontend/onnx/onnx_frontend.py", line 1154, in _impl_v1
    raise ValueError(f"Unsupported PRelu slope shape: {slope_shape}")
ValueError: Unsupported PRelu slope shape: R.shape([64, 1, 1])

Additional context

• The ONNX PRelu spec states: "The shape of slope can be smaller than first input X; if so, its shape must be unidirectional broadcastable to X." It does not restrict slope to rank-1 or to a single non-broadcast channel axis.
• A minimal lower-rank case (X: (2, 3, 4, 5), slope: (3, 1, 1) → channel dim) reproduces the same ValueError with onnx_frontend.py:1154.
• A broader differential survey (7 X shapes × every valid slope broadcast shape): of 113 legal ONNX PRelu models accepted by onnxruntime / onnx.checker / onnx.reference, TVM's frontend rejects 67 — lower-rank slope (32), same-rank multi-non-broadcast including slope == X (28), scalar rank-0 slope (7, crashing with IndexError). No numeric mismatch was observed on the 46 accepted models, so the gap is purely import coverage.
• relax.op.nn.prelu can already express any single per-axis slope (the same-rank branch calls nn.prelu(x, slope, axis)), so the lower-rank case is expressible by computing axis = ndim - s_ndim + relative_axis before reshaping the slope, rather than rejecting the model.
• A fix was upstreamed in apache/tvm #20115 "[Relax][ONNX] Support lower-rank PRelu slopes" (merged 2026-08-11), whose motivating case is exactly X(1,64,128,128) + slope(64,1,1) from a Qualcomm Real-ESRGAN export — confirming the bug is real and observed in the wild. This report covers the unfixed behavior in the v0.24.dev0 build at commit 262c6d2e0.

Triage

• needs-triage
• bug
• relax
• frontend/onnx