diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 1c544d309717b..62396d839dc9b 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -822,13 +822,11 @@ class Prelu(OnnxOpConverter): @classmethod def _impl_v1(cls, inputs, attr, params): assert len(inputs) == 2, "Prelu need 2 inputs, {} given".format(len(inputs)) - input_channels = infer_shape(inputs[0])[1] - alpha_shape = infer_shape(inputs[1]) - if len(alpha_shape) != 1: - alpha = _op.reshape(inputs[1], (-1,)) - else: - alpha = inputs[1] - return _op.nn.prelu(inputs[0], _op.broadcast_to(alpha, [input_channels])) + input_shape = _op.shape_of(inputs[0]) + alpha = _op.broadcast_to_like(inputs[1], inputs[0]) + alpha = _op.reshape(alpha, [-1]) + output = _op.nn.prelu(_op.reshape(inputs[0], [-1]), alpha, axis=0) + return _op.reshape(output, input_shape) class Reciprocal(OnnxOpConverter): diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 3d95a9a83ee31..0f7fda7301cda 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -2037,12 +2037,15 @@ def verify_prelu(x_shape, a_shape): model = helper.make_model(graph, producer_name="prelu_test") - verify_with_ort(model, [x_shape, a_shape], list(x_shape)) + verify_with_ort( + model, [x_shape, a_shape], list(x_shape), use_vm=True, convert_to_static=True + ) verify_prelu([3, 4, 5, 6], [1, 4, 1, 1]) verify_prelu([1, 8, 5, 6], [1, 8, 1, 1]) verify_prelu([2, 12, 16, 16], [1, 12, 1, 1]) verify_prelu([2, 12, 16, 16], [1]) # Test alpha broadcasting. + verify_prelu([3, 1], [3, 1]) # Test non NCHW workload. @tvm.testing.uses_gpu