diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index d529aec461eb..12673c5303b9 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -3444,6 +3444,35 @@ def body_fn(*loop_inputs): return outputs +class LinearRegressor(OnnxOpConverter): + """Operator converter for LinearRegressor.""" + + @classmethod + def _impl_v1(cls, inputs, attr, params): + data = inputs[0] + coefficients = attr.get("coefficients", 0) + data_shape = infer_shape(data) + targets = attr.get("targets", 1) + coefficients = _expr.const(list(coefficients), dtype="float32") + coefficients_shape = infer_shape(coefficients) + + coefficients = _op.reshape(coefficients, (targets, coefficients_shape[0] // targets)) + if coefficients_shape[0] // targets < data_shape[-1]: + data = _op.split(data, [coefficients_shape[0] // targets], -1)[0] + + mm_out = _op.nn.dense(data, coefficients) + + if "intercepts" in attr: + intercepts = attr.get("intercepts", 0) + intercepts = _expr.const(list(intercepts), dtype="float32") + + if targets == 1: + return _op.nn.bias_add(mm_out, intercepts, axis=-1) + return get_relay_op("add")(mm_out, intercepts) + + return mm_out + + class NonMaxSuppression(OnnxOpConverter): """Operator converter for NonMaxSuppression.""" @@ -4770,6 +4799,8 @@ def _get_convert_map(opset): "Adam": Adam.get_converter(opset), "Momentum": Momentum.get_converter(opset), "Scan": Scan.get_converter(opset), + # ML + "LinearRegressor": LinearRegressor.get_converter(opset), } diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 84e3ddbf53d2..b62509297300 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -6275,6 +6275,49 @@ def verify_scan( verify_scan(input_shapes, output_shapes, 3, [-4, -1, -2], [1] * 3, [-3, -2], [1] * 2, 9) +@tvm.testing.parametrize_targets +def test_LinearRegressor(target, dev): + def verify_LinearRegressor(a_shape, c_shape, i_shape, targets=1, batch=1): + a_array = np.random.uniform(size=a_shape).astype("float32") + out_shape = (batch, targets) + + coefficients = np.random.uniform(size=c_shape).astype("float32") + intercepts = np.random.uniform(size=i_shape).astype("float32") + + mul_node = helper.make_node( + "LinearRegressor", + ["a"], + ["out"], + coefficients=coefficients, + intercepts=intercepts, + targets=targets, + domain="ai.onnx.ml", + ) + + graph = helper.make_graph( + [mul_node], + "LinearRegressor_test", + inputs=[ + helper.make_tensor_value_info("a", TensorProto.FLOAT, list(a_shape)), + ], + outputs=[helper.make_tensor_value_info("out", TensorProto.FLOAT, out_shape)], + ) + model = helper.make_model( + graph, + producer_name="LinearRegressor_test", + opset_imports=[ + onnx.helper.make_opsetid("ai.onnx.ml", 1), + ], + ) + verify_with_ort_with_inputs(model, [a_array], target=target, dev=dev) + + verify_LinearRegressor((1, 3), (3), (1)) + verify_LinearRegressor((2, 10), (10), (1), batch=2) + verify_LinearRegressor((1, 3), (30), (10), targets=10) + verify_LinearRegressor((10, 3), (30), (10), targets=10, batch=10) + verify_LinearRegressor((1, 4), (3), (1)) + + if __name__ == "__main__": test_flatten() test_reshape() @@ -6371,3 +6414,4 @@ def verify_scan( test_random_uniform_like() test_random_normal() test_random_normal_like() + test_LinearRegressor()