|
|
|
@ -16,7 +16,6 @@
|
|
|
|
|
import mindspore as ms
|
|
|
|
|
from mindspore.ops import prim_attr_register, PrimitiveWithInfer
|
|
|
|
|
from mindspore.ops.composite import multitype_ops as C
|
|
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
|
__all__ = ["CusBatchMatMul",
|
|
|
|
|
"CusCholeskyTrsm",
|
|
|
|
@ -31,6 +30,7 @@ __all__ = ["CusBatchMatMul",
|
|
|
|
|
"CusMatMulCubeFraczLeftCast",
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class CusBatchMatMul(PrimitiveWithInfer):
|
|
|
|
|
"""CusBatchMatMul definition"""
|
|
|
|
|
|
|
|
|
@ -38,7 +38,7 @@ class CusBatchMatMul(PrimitiveWithInfer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusBatchMatMul"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matrix_combine_impl import CusMatrixCombine
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x1, x2, out, dout):
|
|
|
|
|
return (C.zeros_like(x1), C.zeros_like(x2))
|
|
|
|
@ -59,7 +59,7 @@ class CusCholeskyTrsm(PrimitiveWithInfer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusCholeskyTrsm"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.cholesky_trsm_impl import CusCholeskyTrsm
|
|
|
|
|
|
|
|
|
|
def infer_shape(self, data1_shape):
|
|
|
|
|
ll = []
|
|
|
|
|
m, _ = data1_shape
|
|
|
|
@ -81,7 +81,7 @@ class CusFusedAbsMax1(PrimitiveWithInfer):
|
|
|
|
|
"""init CusFusedAbsMax1"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1'], outputs=['y'])
|
|
|
|
|
self.origin_shape = origin_shape
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.fused_abs_max1_impl import CusFusedAbsMax1
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x, out, dout):
|
|
|
|
|
return (C.zeros_like(x),)
|
|
|
|
@ -91,7 +91,7 @@ class CusFusedAbsMax1(PrimitiveWithInfer):
|
|
|
|
|
def infer_shape(self, data1_shape):
|
|
|
|
|
ll = []
|
|
|
|
|
if len(data1_shape) == 2:
|
|
|
|
|
ll = [1,]
|
|
|
|
|
ll = [1, ]
|
|
|
|
|
else:
|
|
|
|
|
ll = [32, 64]
|
|
|
|
|
return ll
|
|
|
|
@ -111,7 +111,7 @@ class CusImg2Col(PrimitiveWithInfer):
|
|
|
|
|
self.strides = strides
|
|
|
|
|
self.dilates = dilates
|
|
|
|
|
self.mode = mode
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.img2col_impl import CusImg2Col
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x, out, dout):
|
|
|
|
|
return (C.zeros_like(x),)
|
|
|
|
@ -141,7 +141,7 @@ class CusMatMulCubeDenseLeft(PrimitiveWithInfer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusMatMulCubeDenseLeft"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matmul_cube_dense_left_impl import CusMatMulCubeDenseLeft
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x1, x2, out, dout):
|
|
|
|
|
return (C.zeros_like(x1), C.zeros_like(x2))
|
|
|
|
@ -162,7 +162,7 @@ class CusMatMulCubeFraczRightMul(PrimitiveWithInfer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusMatMulCubeFraczRightMul"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2', 'x3'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matmul_cube_fracz_right_mul_impl import CusMatMulCubeFraczRightMul
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x1, x2, x3, out, dout):
|
|
|
|
|
return (C.zeros_like(x1), C.zeros_like(x2), C.zeros_like(x3))
|
|
|
|
@ -185,7 +185,7 @@ class CusMatMulCube(PrimitiveWithInfer):
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2'], outputs=['y'])
|
|
|
|
|
self.transpose_a = transpose_a
|
|
|
|
|
self.transpose_b = transpose_b
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matmul_cube_impl import CusMatMulCube
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x1, x2, out, dout):
|
|
|
|
|
return (C.zeros_like(x1), C.zeros_like(x2))
|
|
|
|
@ -218,7 +218,7 @@ class CusMatrixCombine(PrimitiveWithInfer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusMatrixCombine"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matrix_combine_impl import CusMatrixCombine
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x, out, dout):
|
|
|
|
|
return (C.zeros_like(x),)
|
|
|
|
@ -242,7 +242,7 @@ class CusTranspose02314(PrimitiveWithInfer):
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusTranspose02314"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.transpose_02314_impl import CusTranspose02314
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x, out, dout):
|
|
|
|
|
return (C.zeros_like(x),)
|
|
|
|
@ -260,36 +260,40 @@ class CusTranspose02314(PrimitiveWithInfer):
|
|
|
|
|
def infer_dtype(self, data1_dtype):
|
|
|
|
|
return data1_dtype
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class CusMatMulCubeDenseRight(PrimitiveWithInfer):
|
|
|
|
|
"""CusMatMulCubeDenseRight definition"""
|
|
|
|
|
|
|
|
|
|
@prim_attr_register
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusMatMulCubeDenseRight"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2','x3'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matmul_cube_dense_right_impl import CusMatMulCubeDenseRight
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2', 'x3'], outputs=['y'])
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x1, x2, x3, out, dout):
|
|
|
|
|
return (C.zeros_like(x1),C.zeros_like(x2), C.zeros_like(x3))
|
|
|
|
|
return (C.zeros_like(x1), C.zeros_like(x2), C.zeros_like(x3))
|
|
|
|
|
|
|
|
|
|
return bprop
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def infer_shape(self, data1_shape, data2_shape, data3_shape):
|
|
|
|
|
return data1_shape
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def infer_dtype(self, data1_dtype, data2_dtype, data3_dtype):
|
|
|
|
|
return ms.common.dtype.tensor_type(getattr(ms, "float32"))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class CusMatMulCubeFraczLeftCast(PrimitiveWithInfer):
|
|
|
|
|
"""CusMatMulCubeFraczLeftCast definition"""
|
|
|
|
|
|
|
|
|
|
@prim_attr_register
|
|
|
|
|
def __init__(self):
|
|
|
|
|
"""init CusMatMulCubeFraczLeftCast"""
|
|
|
|
|
self.init_prim_io_names(inputs=['x1', 'x2'], outputs=['y'])
|
|
|
|
|
from mindspore.ops._op_impl._custom_op.matmul_cube_fracz_left_cast_impl import CusMatMulCubeFraczLeftCast
|
|
|
|
|
|
|
|
|
|
def get_bprop(self):
|
|
|
|
|
def bprop(x1, x2, out, dout):
|
|
|
|
|
return (C.zeros_like(x1),C.zeros_like(x2))
|
|
|
|
|
return (C.zeros_like(x1), C.zeros_like(x2))
|
|
|
|
|
|
|
|
|
|
return bprop
|
|
|
|
|
|
|
|
|
|
def infer_shape(self, data1_shape, data2_shape):
|
|
|
|
|