未验证 提交 ce8455c0 编写于 作者: H HongyuJia 提交者: GitHub

[0D-Tensor] Support matmul, fix infershape (#55316)

上级 3559252a
......@@ -90,7 +90,8 @@ std::vector<std::vector<int>> GetMatmulNewShapes(
: std::vector<int>{1, x_shape[0]};
new_y_shape = trans_y ? std::vector<int>{1, y_shape[0]}
: std::vector<int>{y_shape[0], 1};
out_shape = {1};
// [m] * [m] -> [], which aligns with Paddle's matmul
out_shape = {};
} else if (x_dim == 1) {
// vector * matrix
int y_K = trans_y ? y_shape[max_dim - 1] : y_shape[max_dim - 2];
......
......@@ -824,5 +824,50 @@ class TestSqueezeOp2D(TestSqueezeOp):
self.target_shape = ()
@OpTestTool.skip_if(
not is_compiled_with_cuda(), "x86 test will be skipped due to timeout."
)
class TestMatmulOp(OpTest):
def setUp(self):
np.random.seed(2023)
self.dtype = "float32"
self.init_input()
def init_input(self):
self.inputs = {
"x": np.random.randint(-10, 10, [10]).astype(self.dtype),
"y": np.random.randint(-10, 10, [10]).astype(self.dtype),
}
self.target_shape = ()
def build_paddle_program(self, target):
x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
y = paddle.to_tensor(self.inputs["y"], stop_gradient=False)
out = paddle.matmul(x, y)
self.paddle_outputs = [out]
def build_cinn_program(self, target):
builder = NetBuilder("matmul_op")
x = builder.create_input(
cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x"
)
y = builder.create_input(
cinn_dtype_convert(self.dtype), self.inputs["y"].shape, "y"
)
out = builder.matmul(x, y)
prog = builder.build()
res = self.get_cinn_output(
prog, target, [x, y], [self.inputs["x"], self.inputs["y"]], [out]
)
self.cinn_outputs = res
self.assertEqual(res[0].shape, self.target_shape)
def test_check_results(self):
self.check_outputs_and_grads()
if __name__ == "__main__":
unittest.main()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册