From fd05ca2de6bda321f0b882c42d080472350227b9 Mon Sep 17 00:00:00 2001 From: wangzhuo325 Date: Mon, 22 Jun 2020 22:19:12 +0800 Subject: [PATCH] call matmul instead of batchmatmul --- python/akg/ms/cce/matmul.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/akg/ms/cce/matmul.py b/python/akg/ms/cce/matmul.py index 9e4624a..ab090ab 100644 --- a/python/akg/ms/cce/matmul.py +++ b/python/akg/ms/cce/matmul.py @@ -15,8 +15,10 @@ # limitations under the License. """matmul""" -from akg.ops.nn import batchmatmul +from akg.ops.nn import matmul -def MatMul(x1, x2, transpose_a=False, transpose_b=False): +def MatMul(x1, x2, out_dtype, transpose_a=False, transpose_b=False): """matmul""" - return batchmatmul.batchmatmul(x1, x2, transpose_a, transpose_b) + return matmul.matmul(x=x1, y=x2, b=None, out_dtype=out_dtype, + left_format="zN", right_format="zN", out_format="zN", + transpose_x=transpose_a, transpose_y=transpose_b) -- GitLab