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}") diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 4adc3a0ab3d4..baf7d49d8e12 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -3554,12 +3554,28 @@ def main( 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(): @@ -3586,6 +3602,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)