mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -25,7 +25,7 @@ from mindspore.ops.functional import broadcast_gradient_args
 from mindspore.ops import operations as P
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.ops.operations.math_ops import Trace, Bernoulli, Renorm
-from mindspore import nn, ops, Tensor
+from mindspore import nn, Tensor
 from mindspore.ops.operations.math_ops import Real, Imag, Complex, Angle
 from mindspore.ops.operations.math_ops import Polar
 from mindspore.ops.operations.math_ops import ComplexAbs
@@ -60,6 +60,7 @@ from mindspore.ops.operations.math_ops import NanToNum
 from mindspore.ops.operations.math_ops import FFTWithSize
 from mindspore.ops.operations.math_ops import Betainc
 from mindspore.ops.operations.math_ops import Cholesky
+from mindspore.ops.operations.math_ops import Fmin
 from mindspore.ops.operations.math_ops import CholeskySolve
 from mindspore.ops.operations.math_ops import InplaceIndexAdd
 from mindspore.ops.operations.math_ops import AddV2
@@ -67,11 +68,12 @@ from mindspore.ops.operations.math_ops import TridiagonalMatMul
 from mindspore.ops.operations.math_ops import TridiagonalSolve
 from mindspore.ops.operations.math_ops import Logit
 from mindspore.ops.operations.math_ops import Diagonal
+from mindspore.ops.operations.math_ops import EuclideanNorm
 from mindspore.ops.operations.array_ops import Transpose, MatrixSetDiagV3
+from mindspore.ops.operations.math_ops import Fmax
 from mindspore.ops.operations._inner_ops import DynamicBroadcastGradientArgs
 from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
-from mindspore.ops.primitive import constexpr
-from mindspore.ops._utils.utils import is_shape_unknown, is_dim_unknown
+from mindspore.ops.primitive import _primexpr
 from mindspore.ops._grad.grad_base import bprop_getters, create_tensor_by_element, dyn_rank
 from mindspore.ops._grad.grad_base import dyn_ones, dyn_fill, sum_grad_reduce_axis
 from mindspore.ops._grad.grad_math_ops import binop_grad_common
@@ -83,7 +85,7 @@ dyn_shape_op = P.TensorShape()
 _conj = P.Conj()
 
 
-@constexpr
+@_primexpr
 def _generate_perm(x_dim):
     perm = tuple(range(x_dim - 2))
     return perm
@@ -158,7 +160,7 @@ def get_bprop_cdist(self):
 
     def bprop(input_x, input_y, out, dout):
         dout_shape = F.shape(dout)
-        if is_shape_unknown(dout_shape):
+        if F.is_sequence_value_unknown(dout_shape):
             dout_dim = dyn_rank(dout)
             dout_perm_part2 = create_tensor_by_element(
                 (dout_dim - 1, dout_dim - 2))
@@ -356,7 +358,7 @@ def get_bprop_index_addcmul(self):
     return bprop
 
 
-@constexpr
+@_primexpr
 def renew_dim(shape, dim):
     """ Re-new dims"""
     new_dim = dim if dim >= 0 else len(shape) + dim
@@ -365,6 +367,21 @@ def renew_dim(shape, dim):
     return tuple(tmp)
 
 
+@bprop_getters.register(EuclideanNorm)
+def get_bprop_euclidean_norm(self):
+    """Generate bprop for EuclideanNorm"""
+    expand_dims = P.ExpandDims()
+    keep_dims = self.keep_dims
+
+    def bprop(x, axes, out, dout):
+        scale_v = dout / out
+        if not keep_dims and x.shape != ():
+            scale_v = expand_dims(scale_v, axes)
+        return (x * scale_v, zeros_like(axes))
+
+    return bprop
+
+
 @bprop_getters.register(Renorm)
 def get_bprop_renorm(self):
     """Generate bprop for Renorm """
@@ -437,7 +454,7 @@ def get_bprop_lp_norm(self):
         else:
             input_scaled = pow_op(abs_op(input_x), (p - 2)) * input_x
             scale_v = dout / pow_op(out, (p - 1))
-        return (input_scaled * scale_v,)
+        return (mnp.where(input_scaled == 0, 0, input_scaled * scale_v),)
 
     return bprop
 
@@ -451,42 +468,18 @@ def get_brop_cumulative_logsumexp(self):
     cumulative_op = CumulativeLogsumexp(self.exclusive, not self.reverse)
     less_op = P.Less()
     neg_op = P.Neg()
-
-    def where_v2(condition, x=None, y=None):
-        return_all = None
-        if x is None and y is None:
-            return_all = mnp.where(condition, x, y)
-        elif x is not None and y is not None:
-            shape_ = x.shape
-            input_y = np.resize(y, shape_)
-            input_y = Tensor(input_y).astype(x.dtype)
-            return_all = ops.select(condition, x, input_y)
-        else:
-            raise ValueError("x and y must both be non-None or both be None.")
-        return return_all
+    cast = P.Cast()
 
     def bprop(x, axis, out, dout):
         dtype_min = 0
-        fp64_flag = False
         if x.dtype == mstype.float16:
-            dtype_min = -65500e+0
-        elif x.dtype == mstype.float32:
-            dtype_min = -3.4028235e+38
-        elif x.dtype == mstype.float64:
-            dout = F.cast(dout, mstype.float32)
-            x = F.cast(x, mstype.float32)
-            out = F.cast(out, mstype.float32)
-            dtype_min = -3.4028235e+38
-            fp64_flag = True
-
-        log_grad_positive = where_v2(greater_op(dout, 0), log_op(dout), dtype_min)
-        log_grad_negative = where_v2(less_op(dout, 0), log_op(neg_op(dout)), dtype_min)
+            dtype_min = cast(np.finfo(np.float16).min, x.dtype)
+        else:
+            dtype_min = cast(np.finfo(np.float32).min, x.dtype)
+        log_grad_positive = mnp.where(greater_op(dout, 0), log_op(dout), dtype_min)
+        log_grad_negative = mnp.where(less_op(dout, 0), log_op(neg_op(dout)), dtype_min)
         output_pos = exp_op(cumulative_op(log_grad_positive - out, axis) + x)
         output_neg = exp_op(cumulative_op(log_grad_negative - out, axis) + x)
-        if fp64_flag:
-            output_pos = F.cast(output_pos, mstype.float64)
-            output_neg = F.cast(output_neg, mstype.float64)
-            x = F.cast(x, mstype.float64)
         return (output_pos - output_neg, zeros_like(x))
 
     return bprop
@@ -561,7 +554,7 @@ def get_bprop_matrix_exp(self):
     dyn_shape = P.TensorShape()
 
     def bprop(x, out, dout):
-        if is_shape_unknown(x.shape):
+        if F.is_sequence_value_unknown(x.shape):
             shape_x = dyn_shape(x)
             x_len = dyn_rank(x)
             input_perm = range_op(cast(0, mstype.int64), x_len, cast(1, mstype.int64))
@@ -585,7 +578,7 @@ def get_bprop_matrix_exp(self):
         meta_grad = concat_col((meta_grad_up, meta_grad_down))
         meta_grad = matrix_exp(meta_grad)
 
-        if is_shape_unknown(x.shape):
+        if F.is_sequence_value_unknown(x.shape):
             begins = dyn_fill(mstype.int32, expand_dims(x_len, 0), 0)
             sizes = cast(shape_x, mstype.int32)
         else:
@@ -600,22 +593,6 @@ def get_bprop_matrix_exp(self):
     return bprop
 
 
-@bprop_getters.register(P.MatrixInverse)
-def get_bprop_matrix_inverse(self):
-    """Generate bprop for MatrixInverse"""
-    matmul_x1 = nn.MatMul(transpose_x1=True)
-    matmul_x2 = nn.MatMul(transpose_x2=True)
-    neg = P.Neg()
-
-    def bprop(x, out, dout):
-        dx = matmul_x2(dout, out)
-        dx = matmul_x1(out, dx)
-        dx = neg(dx)
-        return (dx,)
-
-    return bprop
-
-
 @bprop_getters.register(MatrixPower)
 def get_bprop_matrix_power(self):
     """Generate bprop for MatrixPower"""
@@ -672,7 +649,7 @@ def get_bprop_matrix_solve(self):
             grad_b = cast(grad_b, mstype.float32)
 
         a_shape = F.shape(input_a)
-        if is_shape_unknown(a_shape):
+        if F.is_sequence_value_unknown(a_shape):
             matrix_rank = dyn_rank(input_a)
         else:
             matrix_rank = rank(input_a)
@@ -697,7 +674,7 @@ def get_bprop_matrix_solve(self):
     return bprop
 
 
-@constexpr
+@_primexpr
 def _generate_perm_matrix_solve_ls(x_dim):
     perm = tuple(range(x_dim - 2))
     perm = perm + (x_dim-1, x_dim-2)
@@ -722,6 +699,7 @@ def get_bprop_matrix_solve_ls(self):
     conjugate_transpose = ConjugateTranspose()
     shape = P.Shape()
     _complex = P.Complex()
+    scalar2tensor = P.ScalarToTensor()
 
     def regularized_gramian_cholesky(matrix, l2, first_kind):
         matrix_dim = rank(matrix)
@@ -785,7 +763,7 @@ def get_bprop_matrix_solve_ls(self):
                 grad_a = add(neg(matmul(matrix, zx_sym)), matmul(rhs, z_temp))
                 grad_b = matmul(matrix, z)
 
-            return (grad_a, grad_b, None)
+            return (grad_a, grad_b, scalar2tensor(0, l2.dtype))
 
         def under_determined(matrix, rhs, l2, dout):
             if matrix.dtype == mstype.complex64:
@@ -832,7 +810,7 @@ def get_bprop_matrix_solve_ls(self):
                 a2 = matmul(tmp, a2_temp)
 
             grad_a = add(a1, a2)
-            return (grad_a, grad_b, None)
+            return (grad_a, grad_b, scalar2tensor(0, l2.dtype))
 
         if fast is False:
             raise ValueError("For MatrixSolveLs, gradient not defined for fast=False")
@@ -855,7 +833,7 @@ def get_bprop_matrix_determinant(self):
     concat = P.Concat(0)
 
     def bprop(x, out, dout):
-        if is_shape_unknown(shape_op(x)):
+        if F.is_sequence_value_unknown(shape_op(x)):
             x_adj_inv = inverse_op(x)
             out_shape = dyn_shape_op(out)
             ones = create_tensor_by_element((1, 1))
@@ -879,7 +857,7 @@ def get_bprop_log_matrix_determinant(self):
 
     def bprop(x, out, dout):
         x_adj_inv = inverse_op(x)
-        if is_shape_unknown(shape_op(out[1])):
+        if F.is_sequence_value_unknown(shape_op(out[1])):
             const_value = F.cast(1, mstype.int64)
             const_value = P.ExpandDims()(const_value, 0)
             new_shape = P.Concat()((dyn_shape_op(out[1]), const_value, const_value))
@@ -902,7 +880,7 @@ def get_bprop_betainc(self):
     dyn_shape = P.TensorShape()
 
     def bprop(input_a, input_b, input_x, out, dout):
-        if is_shape_unknown(F.shape(input_x)):
+        if F.is_sequence_value_unknown(F.shape(input_x)):
             sx = dyn_shape(input_x)
         else:
             sx = F.shape(input_x)
@@ -1044,6 +1022,7 @@ def get_bprop_polar(self):
     sig = P.Sign()
     ones = P.Ones()
     zeros = P.Zeros()
+
     def bprop(input1, angle, out, dout):
         grad_conj = conj(dout)
         zero = zeros(dout.shape, input1.dtype)
@@ -1317,7 +1296,7 @@ def get_bprop_trace(self):
 
     def bprop(x, out, dout):
         shape = shape_op(x)
-        if is_shape_unknown(shape):
+        if F.is_sequence_value_unknown(shape):
             shape = dyn_shape_op(x)
             dx = input_grad(dout, shape)
         else:
@@ -1327,6 +1306,115 @@ def get_bprop_trace(self):
     return bprop
 
 
+@bprop_getters.register(Fmin)
+def get_bprop_fmin(self):
+    """Grad definition for 'Fmin' operation"""
+    shape_ = P.Shape()
+    masked_fill_op = P.MaskedFill()
+    logical_or_op = P.LogicalOr()
+    logical_not_op = P.LogicalNot()
+    logical_and_op = P.LogicalAnd()
+    mul_op = P.Mul()
+    is_nan_op = P.IsNan()
+    reshape_ = P.Reshape()
+
+    def bprop(x1, x2, out, dout):
+        x1_dtype = F.dtype(x1)
+        x2_dtype = F.dtype(x2)
+        x1 = F.cast(x1, mstype.float32)
+        x2 = F.cast(x2, mstype.float32)
+        dout = F.cast(dout, mstype.float32)
+        b1 = logical_or_op((x1 <= x2), is_nan_op(x2))
+        b2 = logical_or_op((x2 < x1), logical_and_op(is_nan_op(x1), logical_not_op(is_nan_op(x2))))
+        rx1 = masked_fill_op(x1, b1, 1.)
+        rx1 = masked_fill_op(rx1, logical_not_op(b1), 0.)
+        rx2 = masked_fill_op(x2, b2, 1.)
+        rx2 = masked_fill_op(rx2, logical_not_op(b2), 0.)
+        rrx1 = mul_op(rx1, dout)
+        rrx2 = mul_op(rx2, dout)
+        shape_of_x1 = shape_(x1)
+        shape_of_x2 = shape_(x2)
+        x1_dim = len(shape_of_x1)
+        x2_dim = len(shape_of_x2)
+        if x1_dim == 0 and x2_dim != 0:
+            sum_r1 = rrx1.sum()
+            sum_r2 = rrx2
+        elif x1_dim == 0 and x2_dim == 0:
+            sum_r1 = rrx1.sum()
+            sum_r2 = rrx2.sum()
+        elif x1_dim != 0 and x2_dim == 0:
+            sum_r2 = rrx2.sum()
+            sum_r1 = rrx1
+        else:
+            rx, ry = DynamicBroadcastGradientArgs()(shape_of_x1, shape_of_x2)
+            sum_r1 = sum_grad_reduce_axis(rrx1, rx)
+            sum_r2 = sum_grad_reduce_axis(rrx2, ry)
+        brrx1 = reshape_(sum_r1, shape_of_x1)
+        brrx2 = reshape_(sum_r2, shape_of_x2)
+        brrx1 = F.cast(brrx1, x1_dtype)
+        brrx2 = F.cast(brrx2, x2_dtype)
+        return brrx1, brrx2
+
+    return bprop
+
+
+@bprop_getters.register(Fmax)
+def get_bprop_fmax(self):
+    """Grad definition for 'Fmax' operation"""
+    shape_ = P.Shape()
+    masked_fill_op = P.MaskedFill()
+    logical_or_op = P.LogicalOr()
+    logical_not_op = P.LogicalNot()
+    logical_and_op = P.LogicalAnd()
+    mul_op = P.Mul()
+    is_nan_op = P.IsNan()
+    reshape_ = P.Reshape()
+
+    def bprop(x1, x2, out, dout):
+        x1_dtype = F.dtype(x1)
+        x2_dtype = F.dtype(x2)
+        if x1_dtype != mstype.float32:
+            x1 = F.cast(x1, mstype.float32)
+            dout = F.cast(dout, mstype.float32)
+        if x2_dtype != mstype.float32:
+            x2 = F.cast(x2, mstype.float32)
+            dout = F.cast(dout, mstype.float32)
+        b1 = logical_or_op(logical_and_op((x1 >= x2), logical_not_op(is_nan_op(x1))), is_nan_op(x2))
+        b2 = logical_or_op(logical_and_op(x2 > x1, logical_not_op(is_nan_op(x2))),
+                           logical_and_op(is_nan_op(x1), logical_not_op(is_nan_op(x2))))
+        rx1 = masked_fill_op(x1, b1, 1.)
+        rx1 = masked_fill_op(rx1, logical_not_op(b1), 0.)
+        rx2 = masked_fill_op(x2, b2, 1.)
+        rx2 = masked_fill_op(rx2, logical_not_op(b2), 0.)
+        rrx1 = mul_op(rx1, dout)
+        rrx2 = mul_op(rx2, dout)
+        shape_of_x1 = shape_(x1)
+        shape_of_x2 = shape_(x2)
+        x1_dim = len(shape_of_x1)
+        x2_dim = len(shape_of_x2)
+        if x1_dim == 0 and x2_dim != 0:
+            sum_r1 = rrx1.sum()
+            sum_r2 = rrx2
+        elif x1_dim == 0 and x2_dim == 0:
+            sum_r1 = rrx1.sum()
+            sum_r2 = rrx2.sum()
+        elif x1_dim != 0 and x2_dim == 0:
+            sum_r2 = rrx2.sum()
+            sum_r1 = rrx1
+        else:
+            rx, ry = DynamicBroadcastGradientArgs()(shape_of_x1, shape_of_x2)
+            sum_r1 = sum_grad_reduce_axis(rrx1, rx)
+            sum_r2 = sum_grad_reduce_axis(rrx2, ry)
+        brrx1 = reshape_(sum_r1, shape_of_x1)
+        brrx2 = reshape_(sum_r2, shape_of_x2)
+        brrx1 = F.cast(brrx1, x1_dtype)
+        brrx2 = F.cast(brrx2, x2_dtype)
+        return brrx1, brrx2
+
+
+    return bprop
+
+
 @bprop_getters.register(G.MinimumGrad)
 def get_bprop_minimum_grad(self):
     """Grad definition for 'MinimumGrad' operation"""
@@ -1355,6 +1443,7 @@ def get_bprop_bernoulli(self):
 def get_bprop_tridiagonalsolve(self):
     """Grad definition for 'TridiagonalSolve' operation"""
     tridiagonalsolve = TridiagonalSolve()
+
     def bprop(diagonals, rhs, out, dout):
         diags = diagonals
         diag1 = diags[..., 1, :]
@@ -1388,7 +1477,7 @@ def get_bprop_igamma(self):
     def bprop(a, x, out, dout):
         sa = shape_(a)
         sx = shape_(x)
-        if is_shape_unknown(sa) or is_shape_unknown(sx):
+        if F.is_sequence_value_unknown(sa) or F.is_sequence_value_unknown(sx):
             sa = dyn_shape_op(a)
             sx = dyn_shape_op(x)
             ra, rx = DynamicBroadcastGradientArgs()(sa, sx)
@@ -1428,7 +1517,7 @@ def get_bprop_igammac(self):
     def bprop(a, x, out, dout):
         sa = shape_(a)
         sx = shape_(x)
-        if is_shape_unknown(sa) or is_shape_unknown(sx):
+        if F.is_sequence_value_unknown(sa) or F.is_sequence_value_unknown(sx):
             sa = dyn_shape_op(a)
             sx = dyn_shape_op(x)
             ra, rx = DynamicBroadcastGradientArgs()(sa, sx)
@@ -1463,6 +1552,9 @@ def get_bprop_lgamma(self):
             x = F.cast(x, mstype.float32)
             dx = dout * digamma(x)
             dx = F.cast(dx, mstype.float16)
+        elif x.dtype in (mstype.int32,):
+            x = F.cast(x, mstype.float32)
+            dx = dout * digamma(x)
         else:
             dx = dout * digamma(x)
         return (dx,)
@@ -1474,13 +1566,15 @@ def get_bprop_lgamma(self):
 def get_bprop_digamma(self):
     """Grad definition for `Digamma` operation."""
     polygamma = Polygamma()
+    a = Tensor(1)
+
     def bprop(x, out, dout):
         if x.dtype in (mstype.float16,):
             x = F.cast(x, mstype.float32)
-            dx = dout * polygamma(1, x)
+            dx = dout * polygamma(a, x)
             dx = F.cast(dx, mstype.float16)
         else:
-            dx = dout * polygamma(1, x)
+            dx = dout * polygamma(a, x)
         return (dx,)
 
     return bprop
@@ -1492,14 +1586,15 @@ def get_bprop_polygamma(self):
     polygamma = Polygamma()
 
     def bprop(a, x, out, dout):
-        a = a + 1
+        one = Tensor(1)
+        a = a + one
         if x.dtype in (mstype.float16,):
             x = F.cast(x, mstype.float64)
             dx = dout * polygamma(a, x)
             dx = F.cast(dx, mstype.float16)
         else:
             dx = dout * polygamma(a, x)
-        return None, dx
+        return zeros_like(a), dx
 
     return bprop
 
@@ -1584,7 +1679,7 @@ def get_bprop_cholesky_solve(self):
     def bprop(x1, x2, out, dout):
         flag = 0
         shape_x1 = shape_op(x1)
-        if is_dim_unknown(shape_x1):
+        if F.is_sequence_shape_unknown(shape_x1):
             len_x1 = dyn_rank(x1)
         else:
             len_x1 = len(shape_x1)
@@ -1643,7 +1738,7 @@ def get_bprop_nextafter(self):
 
         s_x1 = shape(x1)
         partial_x1 = ()
-        if is_shape_unknown(s_x1):
+        if F.is_sequence_value_unknown(s_x1):
             s_x1 = dyn_shape(x1)
             partial_x1 = dyn_ones(s_x1, dtype(x1))
         else:
@@ -1651,7 +1746,7 @@ def get_bprop_nextafter(self):
 
         s_x2 = shape(x2)
         partial_x2 = ()
-        if is_shape_unknown(s_x2):
+        if F.is_sequence_value_unknown(s_x2):
             s_x2 = dyn_shape(x2)
             partial_x2 = dyn_fill(dtype(x2), s_x2, 0)
         else:
@@ -1670,13 +1765,14 @@ def get_bprop_diagonal(self):
     offset = self.offset
     dim1 = self.dim1
     dim2 = self.dim2
-    zeros = P.Zeros()
+    zeros_op = P.FillV2()
     size_op = P.Size()
     transpose_op = Transpose()
     matrix_set_diag_op = MatrixSetDiagV3(align="LEFT_RIGHT")
 
     def bprop(x, out, dout):
         x_shape = x.shape
+        x_dtype = x.dtype
         x_dim = len(x_shape)
         if dim1 < 0:
             dim1_ = dim1 + x_dim
@@ -1690,7 +1786,8 @@ def get_bprop_diagonal(self):
             batch_dim = out.shape[:-1]
             diag_plane = (x_shape[dim1_], x_shape[dim2_])
             dx_trans_shape = batch_dim + diag_plane
-            dx = zeros(dx_trans_shape, x.dtype)
+            value = Tensor(0, x_dtype)
+            dx = zeros_op(dx_trans_shape, value)
             k = F.cast(offset, mstype.int32)
             dx = matrix_set_diag_op(dx, dout, k)
             dim = 0
@@ -1738,13 +1835,28 @@ def get_bprop_inplace_index_add(self):
     return bprop
 
 
-@constexpr
+@bprop_getters.register(P.Zeta)
+def get_bprop_zeta(self):
+    """Generate bprop for Zeta"""
+    zeta = P.Zeta()
+
+    def bprop(x, q, out, dout):
+        dq = -x * zeta(x + 1, q) * dout
+        return zeros_like(x), dq
+
+    return bprop
+
+
+@_primexpr
 def _fft_rank_offset(norm_shape, rank):
     """generate offset for fft with rank"""
-    return np.product(norm_shape[-rank:])
+    norm_shape_product = 1
+    for i in norm_shape[-rank:]:
+        norm_shape_product *= i
+    return norm_shape_product
 
 
-@constexpr
+@_primexpr
 def _fft_with_size_back_norm(norm_shape, norm, inverse, rank):
     """generate reverse term for fft_with_size"""
     if inverse is False:
@@ -1764,7 +1876,7 @@ def _fft_with_size_back_norm(norm_shape, norm, inverse, rank):
     return norm_
 
 
-@constexpr
+@_primexpr
 def _rfft_norm(norm_shape, norm, rank):
     """generate norm for rfft"""
     norm_ = 1.0
@@ -1777,7 +1889,7 @@ def _rfft_norm(norm_shape, norm, rank):
     return norm_
 
 
-@constexpr
+@_primexpr
 def _get_last_dim_slice_shape(tensor_shape, index):
     """generate shape for slice last tensor"""
     from_shape = [0 for x in tensor_shape]
@@ -1790,31 +1902,40 @@ def _get_last_dim_slice_shape(tensor_shape, index):
     return tuple(from_shape), tuple(to_shape)
 
 
-@constexpr
+@_primexpr
 def _rfft_reshape(shape_a, shape_b):
     """generate rfft shape for reshape"""
-    return tuple([int(x) for x in np.concatenate((np.ones(len(shape_a) - 2), shape_b), 0)])
+    new_shape = list(shape_b)
+    for i in range(len(shape_a) - 2):
+        new_shape.insert(i, 1)
+    return tuple(new_shape)
 
 
-@constexpr
+@_primexpr
 def _rfft_tile_reshape(shape_a):
     """generate rfft shape for tile"""
-    return tuple(int(x) for x in np.concatenate((shape_a[:-2], [1, 1]), 0))
+    reshape_a = list(shape_a)
+    reshape_a[-2] = 1
+    reshape_a[-1] = 1
+    return tuple(reshape_a)
 
 
-@constexpr
+@_primexpr
 def _rfft_last_term_shape(shape_a, shape_b):
     """generate rfft shape for last term"""
-    return tuple(int(x) for x in np.concatenate((np.ones(len(shape_a)-1), shape_b), 0))
+    new_shape = list(shape_b)
+    for i in range(len(shape_a) - 1):
+        new_shape.insert(i, 1)
+    return tuple(new_shape)
 
 
-@constexpr
+@_primexpr
 def _batch_matmul_shape_increase(shape_before):
     """increase tensor shape for batch_matmul"""
     return (1, *shape_before)
 
 
-@constexpr
+@_primexpr
 def _batch_matmul_shape_decrease(matrix_shape):
     """decrease tensor shape after batch_matmul"""
     shape_tmp = list(matrix_shape)
@@ -1938,9 +2059,9 @@ def get_bprop_fft_with_size(self):
                             input_shape[-1], input_type))
                         dout_last_term = complex_op(tile_op(real_op(dout_last_term), _rfft_last_term_shape(dout_shape,
                                                                                                            [input_shape[
-                                                                                                               -1],])),
+                                                                                                               -1]])),
                                                     tile_op(imag_op(dout_last_term), _rfft_last_term_shape(
-                                                        dout_shape, [input_shape[-1],])))
+                                                        dout_shape, [input_shape[-1]])))
                         dout_last_term = dout_last_term * vec_mask
                         terms = real_op(dout_first_term) + is_even * real_op(dout_last_term)
                     elif signal_ndim == 3: