mindspore 1.10.0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -17,14 +17,17 @@
 
 import numpy as np
 import mindspore.numpy as mnp
+from mindspore import context
 from mindspore.common import dtype as mstype
-from mindspore import nn
 from mindspore.nn import LGamma
 from mindspore.ops import functional as F
 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, 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
 from mindspore.ops.operations.math_ops import Sinc
 from mindspore.ops.operations import _grad_ops as G
@@ -39,31 +42,50 @@ from mindspore.ops.operations.math_ops import BesselK0e
 from mindspore.ops.operations.math_ops import BesselK1e
 from mindspore.ops.operations.math_ops import BesselY0
 from mindspore.ops.operations.math_ops import BesselY1
+from mindspore.ops.operations.math_ops import Lgamma
+from mindspore.ops.operations.math_ops import Digamma
+from mindspore.ops.operations.math_ops import Polygamma
 from mindspore.ops.operations.math_ops import NextAfter
 from mindspore.ops.operations.math_ops import Hypot
 from mindspore.ops.operations.math_ops import ReduceStd
 from mindspore.ops.operations.math_ops import LuUnpack
 from mindspore.ops.operations.math_ops import MatrixExp
+from mindspore.ops.operations.math_ops import CumulativeLogsumexp
 from mindspore.ops.operations.math_ops import MatrixSolve
+from mindspore.ops.operations.math_ops import MatrixSolveLs
+from mindspore.ops.operations.math_ops import MatrixPower
 from mindspore.ops.operations.math_ops import Median
+from mindspore.ops.operations.math_ops import MatrixTriangularSolve
+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
 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
+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
-
+from mindspore.ops.operations.array_ops import MatrixBandPart
+from mindspore.ops.operations.array_ops import ConjugateTranspose
 
 transpose = P.Transpose()
 dyn_shape_op = P.TensorShape()
 _conj = P.Conj()
 
 
-@constexpr
+@_primexpr
 def _generate_perm(x_dim):
     perm = tuple(range(x_dim - 2))
     return perm
@@ -110,6 +132,27 @@ def get_bprop_logit(self):
     return bprop
 
 
+@bprop_getters.register(P.Roll)
+def get_bprop_roll(self):
+    """Generate bprop for Roll"""
+    if context.get_context("device_target") == "GPU":
+        shift = []
+        axis = self.axis
+        for tmp in enumerate(self.shift):
+            shift.append(-tmp[1])
+        roll_grad = P.Roll(shift, axis)
+    else:
+        shift = self.shift
+        axis = self.axis
+        roll_grad = P.Roll(-shift, axis)
+
+    def bprop(x_input, out, dout):
+        dx = roll_grad(dout)
+        return (dx,)
+
+    return bprop
+
+
 @bprop_getters.register(P.Cdist)
 def get_bprop_cdist(self):
     """Generate bprop for Cdist"""
@@ -117,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))
@@ -146,7 +189,7 @@ def get_bprop_index_lerp(self):
     """Generate bprop for Lerp"""
     mul_op = P.Mul()
     sub_op = P.Sub()
-    is_instance_op = P.IsInstance()
+    is_instance_op = inner.IsInstance()
 
     def bprop(start, end, weight, out, dout):
         dout = F.cast(dout, mstype.float32)
@@ -315,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
@@ -324,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 """
@@ -332,7 +390,6 @@ def get_bprop_renorm(self):
     dim = self.dim
     max_norm = self.maxnorm
     greater_op = P.Greater()
-    masked_fill_op = P.MaskedFill()
     pow_op = P.Pow()
     abs_op = P.Abs()
     sign_op = P.Sign()
@@ -349,13 +406,13 @@ def get_bprop_renorm(self):
             norm_bp = sig * grad_out
         elif p == 2:
             m = input_x * (grad_out / norm)
-            norm_bp = masked_fill_op(m, norm == 0., 0.)
+            norm_bp = F.masked_fill(m, norm == 0., 0.)
         else:
             abs_ = abs_op(input_x)
             input_scaled = input_x * pow_op(abs_, (p - 2))
             pow_ = pow_op(norm, (p - 1))
             scale_v = grad_out / pow_
-            scale_v = masked_fill_op(scale_v, norm == 0., 0.)
+            scale_v = F.masked_fill(scale_v, norm == 0., 0.)
             norm_bp = input_scaled * scale_v
 
         v = norm + ext
@@ -397,7 +454,88 @@ 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
+
+
+@bprop_getters.register(CumulativeLogsumexp)
+def get_brop_cumulative_logsumexp(self):
+    """Generate bprop for CumulativeLogsumexp"""
+    exp_op = P.Exp()
+    greater_op = P.Greater()
+    log_op = P.Log()
+    cumulative_op = CumulativeLogsumexp(self.exclusive, not self.reverse)
+    less_op = P.Less()
+    neg_op = P.Neg()
+    cast = P.Cast()
+
+    def bprop(x, axis, out, dout):
+        dtype_min = 0
+        if x.dtype == mstype.float16:
+            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)
+        return (output_pos - output_neg, zeros_like(x))
+
+    return bprop
+
+
+@bprop_getters.register(MatrixTriangularSolve)
+def get_bprop_matrix_triangular_solve(self):
+    """Grad definition for 'MatrixTriangularSolve' operation"""
+    adjoint_a = self.adjoint
+    lower_a = self.lower
+    matrix_triangular_solve_op = P.MatrixTriangularSolve(lower=lower_a, adjoint=not adjoint_a)
+    mat_mul_2d_op = P.MatMul()
+    mat_mul_op = P.BatchMatMul()
+    real_op = P.Real()
+    imag_op = P.Imag()
+    neg_op = P.Neg()
+    complex_op = P.Complex()
+    matrix_band_part_op = MatrixBandPart()
+
+    def bprop(matrix, rhs, out, dout):
+        grad_rhs = matrix_triangular_solve_op(matrix, dout)
+        if matrix.dtype == mstype.complex64 or matrix.dtype == mstype.complex128:
+            grad_rhs_temp = _adjoint(grad_rhs)
+            out_temp = _adjoint(out)
+        else:
+            grad_rhs_temp = cholesky_transpose(grad_rhs)
+            out_temp = cholesky_transpose(out)
+        if adjoint_a:
+            if len(matrix.shape) == 2:
+                grad_matrix = mat_mul_2d_op(out, grad_rhs_temp)
+                grad_matrix = neg_op(grad_matrix)
+            else:
+                grad_matrix = mat_mul_op(out, grad_rhs_temp)
+                grad_matrix = neg_op(grad_matrix)
+        else:
+            if len(matrix.shape) == 2:
+                grad_matrix = mat_mul_2d_op(grad_rhs, out_temp)
+                grad_matrix = neg_op(grad_matrix)
+            else:
+                grad_matrix = mat_mul_op(grad_rhs, out_temp)
+                grad_matrix = neg_op(grad_matrix)
+        if lower_a:
+            if grad_matrix.dtype == mstype.complex64 or grad_matrix.dtype == mstype.complex128:
+                grad_matrix_real = matrix_band_part_op(real_op(grad_matrix), -1, 0)
+                grad_matrix_imag = matrix_band_part_op(imag_op(grad_matrix), -1, 0)
+                grad_matrix = complex_op(grad_matrix_real, grad_matrix_imag)
+            else:
+                grad_matrix = matrix_band_part_op(grad_matrix, -1, 0)
+        else:
+            if grad_matrix.dtype == mstype.complex64 or grad_matrix.dtype == mstype.complex128:
+                grad_matrix_real = matrix_band_part_op(real_op(grad_matrix), 0, -1)
+                grad_matrix_imag = matrix_band_part_op(imag_op(grad_matrix), 0, -1)
+                grad_matrix = complex_op(grad_matrix_real, grad_matrix_imag)
+            else:
+                grad_matrix = matrix_band_part_op(grad_matrix, 0, -1)
+        return (grad_matrix, grad_rhs)
 
     return bprop
 
@@ -411,27 +549,43 @@ def get_bprop_matrix_exp(self):
     concat_col = P.Concat(-2)
     cast = P.Cast()
     slice_op = P.Slice()
+    range_op = P.Range()
+    expand_dims = P.ExpandDims()
+    dyn_shape = P.TensorShape()
 
     def bprop(x, out, dout):
-        shape_x = P.Shape()(x)
-        n = shape_x[-1]
-        zero_matrix = zeros(shape_x, mstype.float32)
-        zero_matrix = cast(zero_matrix, dout.dtype)
-        x_len = len(shape_x)
-        input_perm = [ele for ele in range(x_len)]
-        input_perm[-1] = input_perm[-2]
-        input_perm[-2] = x_len-1
-        input_perm = tuple(input_perm)
-        x_transpose = P.Transpose()(x, input_perm)
+        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))
+            input_perm[-1] = input_perm[-2]
+            input_perm[-2] = x_len - 1
+            x_transpose = transpose(x, input_perm)
+            zero_matrix = dyn_fill(mstype.float32, shape_x, 0)
+        else:
+            shape_x = x.shape
+            x_len = len(shape_x)
+            input_perm = [ele for ele in range(x_len)]
+            input_perm[-1] = input_perm[-2]
+            input_perm[-2] = x_len - 1
+            input_perm = tuple(input_perm)
+            x_transpose = P.Transpose()(x, input_perm)
+            zero_matrix = zeros(shape_x, mstype.float32)
 
+        zero_matrix = cast(zero_matrix, dout.dtype)
         meta_grad_up = concat_row((x_transpose, dout))
         meta_grad_down = concat_row((zero_matrix, x_transpose))
         meta_grad = concat_col((meta_grad_up, meta_grad_down))
         meta_grad = matrix_exp(meta_grad)
 
-        begins = [0] * x_len
+        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:
+            begins = [0] * x_len
+            sizes = [i for i in shape_x]
+        n = shape_x[-1]
         begins[-1] = n
-        sizes = [i for i in shape_x]
         sizes[-2] = n
         sizes[-1] = n
         return (slice_op(meta_grad, begins, sizes),)
@@ -439,17 +593,35 @@ 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)
+@bprop_getters.register(MatrixPower)
+def get_bprop_matrix_power(self):
+    """Generate bprop for MatrixPower"""
+    n = self.n
+    batch_matmul_a = P.BatchMatMul(transpose_a=True)
+    batch_matmul_b = P.BatchMatMul(transpose_b=True)
     neg = P.Neg()
 
     def bprop(x, out, dout):
-        dx = matmul_x2(dout, out)
-        dx = matmul_x1(out, dx)
-        dx = neg(dx)
+        dout = F.cast(dout, mstype.float32)
+        x = F.cast(x, mstype.float32)
+        power = n
+        dx = zeros_like(x)
+        if power < 0:
+            matrix_power = MatrixPower(n=-1)
+            x_inv = matrix_power(x)
+            for i in range(0, -power):
+                matrix_power = MatrixPower(n=(-power - 1 - i))
+                dx = dx + batch_matmul_b(dout, matrix_power(x_inv))
+                dout = batch_matmul_a(x_inv, dout)
+            dx = batch_matmul_b(dx, x_inv)
+            dx = batch_matmul_a(x_inv, dx)
+            dx = neg(dx)
+        else:
+            for i in range(0, power):
+                matrix_power = MatrixPower(n=(power - 1 - i))
+                dx = dx + batch_matmul_b(dout, matrix_power(x))
+                dout = batch_matmul_a(x, dout)
+        dx = F.cast(dx, F.dtype(out))
         return (dx,)
 
     return bprop
@@ -475,6 +647,13 @@ def get_bprop_matrix_solve(self):
         grad_b_type = F.dtype(grad_b)
         if grad_b_type == mstype.float64:
             grad_b = cast(grad_b, mstype.float32)
+
+        a_shape = F.shape(input_a)
+        if F.is_sequence_value_unknown(a_shape):
+            matrix_rank = dyn_rank(input_a)
+        else:
+            matrix_rank = rank(input_a)
+
         matrix_rank = rank(input_a)
         if adjoint:
             if matrix_rank > 2:
@@ -495,6 +674,156 @@ def get_bprop_matrix_solve(self):
     return bprop
 
 
+@_primexpr
+def _generate_perm_matrix_solve_ls(x_dim):
+    perm = tuple(range(x_dim - 2))
+    perm = perm + (x_dim-1, x_dim-2)
+    return perm
+
+
+@bprop_getters.register(MatrixSolveLs)
+def get_bprop_matrix_solve_ls(self):
+    """Grad definition for 'MatrixSolveLs' operation"""
+    fast = self.fast
+    cast = P.Cast()
+    neg = P.Neg()
+    rank = P.Rank()
+    cholesky = Cholesky()
+    eye = P.Eye()
+    add = P.Add()
+    mul = P.Mul()
+    matmul = P.MatMul()
+    batch_matmul = P.BatchMatMul()
+    cholesky_solve = CholeskySolve()
+    _transpose = Transpose()
+    conjugate_transpose = ConjugateTranspose()
+    shape = P.Shape()
+    _complex = P.Complex()
+    scalar2tensor = P.ScalarToTensor()
+
+    def regularized_gramian_cholesky(matrix, l2, first_kind):
+        matrix_dim = rank(matrix)
+        perm = _generate_perm_matrix_solve_ls(matrix_dim)
+        if matrix.dtype in (mstype.complex64, mstype.complex128):
+            matrix_temp = conjugate_transpose(matrix, perm)
+        else:
+            matrix_temp = _transpose(matrix, perm)
+        if first_kind:
+            if matrix_dim > 2:
+                gramian = batch_matmul(matrix_temp, matrix)
+            else:
+                gramian = matmul(matrix_temp, matrix)
+        else:
+            if matrix_dim > 2:
+                gramian = batch_matmul(matrix, matrix_temp)
+            else:
+                gramian = matmul(matrix, matrix_temp)
+        if isinstance(l2, Tensor) or l2 != 0:
+            matrix_shape = shape(matrix)
+            if first_kind:
+                small_dim = matrix_shape[-1]
+            else:
+                small_dim = matrix_shape[-2]
+            identity = eye(small_dim, small_dim, matrix.dtype)
+            gramian = add(gramian, mul(l2, identity))
+
+        #Cholesky not support complex dtype for now
+        return cholesky(gramian)
+
+    def bprop(matrix, rhs, l2, out, dout):
+        #support dtype:float32
+        #support dimension: 2D,3D
+        def over_determined(matrix, rhs, out, l2, dout):
+            if matrix.dtype == mstype.complex64:
+                l2_regularizer = _complex(cast(l2, mstype.float32), Tensor(0, mstype.float32))
+            elif matrix.dtype == mstype.complex128:
+                l2_regularizer = _complex(cast(l2, mstype.float64), Tensor(0, mstype.float64))
+            else:
+                l2_regularizer = cast(l2, matrix.dtype)
+            chol = cast(regularized_gramian_cholesky(matrix, l2_regularizer, first_kind=True), matrix.dtype)
+            #CholeskySolve not support complex dtype and just support 2D or 3D matrices for now
+            z = cholesky_solve(dout, chol)
+
+            matrix_dim = rank(matrix)
+            perm = _generate_perm_matrix_solve_ls(matrix_dim)
+            if matrix.dtype in (mstype.complex64, mstype.complex128):
+                z_temp = conjugate_transpose(z, perm)
+            else:
+                z_temp = _transpose(z, perm)
+            if matrix_dim > 2:
+                xzt = batch_matmul(out, z_temp)
+            else:
+                xzt = matmul(out, z_temp)
+            zx_sym = add(xzt, _transpose(xzt, perm))
+
+            if matrix_dim > 2:
+                grad_a = add(neg(batch_matmul(matrix, zx_sym)), batch_matmul(rhs, z_temp))
+                grad_b = batch_matmul(matrix, z)
+            else:
+                grad_a = add(neg(matmul(matrix, zx_sym)), matmul(rhs, z_temp))
+                grad_b = matmul(matrix, z)
+
+            return (grad_a, grad_b, scalar2tensor(0, l2.dtype))
+
+        def under_determined(matrix, rhs, l2, dout):
+            if matrix.dtype == mstype.complex64:
+                l2_regularizer = _complex(cast(l2, mstype.float32), Tensor(0, mstype.float32))
+            elif matrix.dtype == mstype.complex128:
+                l2_regularizer = _complex(cast(l2, mstype.float64), Tensor(0, mstype.float64))
+            else:
+                l2_regularizer = cast(l2, matrix.dtype)
+            chol = cast(regularized_gramian_cholesky(matrix, l2_regularizer, first_kind=False), matrix.dtype)
+
+            matrix_dim = rank(matrix)
+            perm = _generate_perm_matrix_solve_ls(matrix_dim)
+            if matrix_dim > 2:
+                gramian = batch_matmul(matrix, dout)
+            else:
+                gramian = matmul(matrix, dout)
+            #CholeskySolve not support complex dtype and just support 2D or 3D matrices for now
+            grad_b = cholesky_solve(gramian, chol)
+            tmp = cholesky_solve(rhs, chol)
+
+            if matrix.dtype in (mstype.complex64, mstype.complex128):
+                tmp_temp = conjugate_transpose(tmp, perm)
+                matrix_temp = conjugate_transpose(matrix, perm)
+            else:
+                tmp_temp = _transpose(tmp, perm)
+                matrix_temp = _transpose(matrix, perm)
+            if matrix_dim > 2:
+                a1 = batch_matmul(tmp_temp, matrix)
+                a1 = neg(batch_matmul(grad_b, a1))
+                a2 = dout - batch_matmul(matrix_temp, grad_b)
+                if matrix.dtype in (mstype.complex64, mstype.complex128):
+                    a2_temp = conjugate_transpose(a2, perm)
+                else:
+                    a2_temp = _transpose(a2, perm)
+                a2 = batch_matmul(tmp, a2_temp)
+            else:
+                a1 = matmul(tmp_temp, matrix)
+                a1 = neg(matmul(grad_b, a1))
+                a2 = dout - matmul(matrix_temp, grad_b)
+                if matrix.dtype in (mstype.complex64, mstype.complex128):
+                    a2_temp = conjugate_transpose(a2, perm)
+                else:
+                    a2_temp = _transpose(a2, perm)
+                a2 = matmul(tmp, a2_temp)
+
+            grad_a = add(a1, a2)
+            return (grad_a, grad_b, scalar2tensor(0, l2.dtype))
+
+        if fast is False:
+            raise ValueError("For MatrixSolveLs, gradient not defined for fast=False")
+        matrix_shape = shape(matrix)[-2:]
+
+        if matrix_shape[-2] >= matrix_shape[-1]:
+            return over_determined(matrix, rhs, out, l2, dout)
+
+        return under_determined(matrix, rhs, l2, dout)
+
+    return bprop
+
+
 @bprop_getters.register(P.MatrixDeterminant)
 def get_bprop_matrix_determinant(self):
     """Generate bprop for MatrixDeterminant"""
@@ -504,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))
@@ -528,7 +857,13 @@ def get_bprop_log_matrix_determinant(self):
 
     def bprop(x, out, dout):
         x_adj_inv = inverse_op(x)
-        multipliers = reshape(dout[1], shape_op(out[1]) + (1, 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))
+            multipliers = reshape(dout[1], new_shape)
+        else:
+            multipliers = reshape(dout[1], shape_op(out[1]) + (1, 1))
         dx = multipliers * x_adj_inv
         return (dx,)
 
@@ -542,18 +877,17 @@ def get_bprop_betainc(self):
     exp = P.Exp()
     log1p = P.Log1p()
     xlogy = P.Xlogy()
-    reduce_sum = P.ReduceSum()
+    dyn_shape = P.TensorShape()
 
     def bprop(input_a, input_b, input_x, out, dout):
-        sa = F.shape(input_a)
-        sx = F.shape(input_x)
-        _, rx = F.broadcast_gradient_args(sa, sx)
-
+        if F.is_sequence_value_unknown(F.shape(input_x)):
+            sx = dyn_shape(input_x)
+        else:
+            sx = F.shape(input_x)
         log_beta = (lgamma(input_a) + lgamma(input_b) - lgamma(input_a + input_b))
         partial_x = exp((input_b - 1) * log1p(-input_x) + xlogy(input_a - 1, input_x) - log_beta)
-        if rx != ():
-            return (zeros_like(input_a), zeros_like(input_b), F.reshape(reduce_sum(partial_x * dout, rx), sx))
         return (zeros_like(input_a), zeros_like(input_b), F.reshape(partial_x * dout, sx))
+
     return bprop
 
 
@@ -646,6 +980,18 @@ def get_bprop_complex_abs(self):
     return bprop
 
 
+@bprop_getters.register(NanToNum)
+def get_bprop_nan_to_num(self):
+    """Grad definition for `NanToNum` operation."""
+    isfinite = P.IsFinite()
+
+    def bprop(x, out, dout):
+        dx = dout * isfinite(x)
+        return (dx,)
+
+    return bprop
+
+
 @bprop_getters.register(Angle)
 def get_bprop_angle(self):
     """Grad definition for `Angle` operation."""
@@ -667,6 +1013,29 @@ def get_bprop_angle(self):
     return bprop
 
 
+@bprop_getters.register(Polar)
+def get_bprop_polar(self):
+    """Grad definition for `Polar` operation."""
+    complex_op = Complex()
+    conj = P.Conj()
+    real = P.Real()
+    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)
+        one = ones(dout.shape, input1.dtype)
+        i = complex_op(zero, one)
+        grad_abs = real(grad_conj * sig(out))
+        result_mul_1_j = out * i
+        grad_angle = real(grad_conj * result_mul_1_j)
+        return (grad_abs, grad_angle)
+
+    return bprop
+
+
 @bprop_getters.register(P.Erfinv)
 def get_bprop_erfinv(self):
     """Grad definition for `Erfinv` operation."""
@@ -927,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:
@@ -937,12 +1306,121 @@ 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"""
     input_grad = G.MinimumGradGrad()
 
-    def bprop(grad, x1, x2, out, dout):
+    def bprop(x1, x2, grad, out, dout):
         sopd_x1, sopd_x2, sopd_grads = input_grad(x1, x2, dout[0], dout[1])
         sopd_x1 = zeros_like(x1)
         sopd_x2 = zeros_like(x2)
@@ -961,6 +1439,30 @@ def get_bprop_bernoulli(self):
     return bprop
 
 
+@bprop_getters.register(TridiagonalSolve)
+def get_bprop_tridiagonalsolve(self):
+    """Grad definition for 'TridiagonalSolve' operation"""
+    tridiagonalsolve = TridiagonalSolve()
+
+    def bprop(diagonals, rhs, out, dout):
+        diags = diagonals
+        diag1 = diags[..., 1, :]
+        zeros1 = P.Zeros()(diags.shape[:-2] + (1,), diags.dtype)
+        superdiag1 = P.Concat(-1)((diags[..., 2, 1:], zeros1))
+        subdiag1 = P.Concat(-1)((zeros1, diags[..., 0, :-1]))
+        diags_transposed = P.Stack(-2)([superdiag1, diag1, subdiag1])
+        grad_rhs = tridiagonalsolve(diags_transposed, dout)
+        diag2 = P.ReduceSum()(grad_rhs * out, -1)
+        zeros2 = P.Zeros()(grad_rhs.shape[:-2] + (1, grad_rhs.shape[-1]), grad_rhs.dtype)
+        superdiag2 = P.ReduceSum()(grad_rhs * P.Concat(-2)((out[..., 1:, :], zeros2)), -1)
+        subdiag2 = P.ReduceSum()(grad_rhs * P.Concat(-2)((zeros2, out[..., :-1, :])), -1)
+        a = (P.Stack(-2)([superdiag2, diag2, subdiag2]))
+        grad_diags = 0 - a
+        return grad_diags, grad_rhs
+
+    return bprop
+
+
 @bprop_getters.register(Igamma)
 def get_bprop_igamma(self):
     """Grad definition for `Igamma` operation."""
@@ -975,6 +1477,15 @@ def get_bprop_igamma(self):
     def bprop(a, x, out, dout):
         sa = shape_(a)
         sx = shape_(x)
+        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)
+            partial_a = igammagrada(a, x)
+            partial_x = exp_(-x + (a - 1) * log_(x) - lgamma(a))
+            r1 = reshape_(sum_grad_reduce_axis(partial_a * dout, ra), sa)
+            r2 = reshape_(sum_grad_reduce_axis(partial_x * dout, rx), sx)
+            return r1, r2
         ra, rx = broadcast_gradient_args(sa, sx)
         partial_a = igammagrada(a, x)
         partial_x = exp_(-x + (a - 1) * log_(x) - lgamma(a))
@@ -1006,6 +1517,15 @@ def get_bprop_igammac(self):
     def bprop(a, x, out, dout):
         sa = shape_(a)
         sx = shape_(x)
+        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)
+            partial_a = igammagrada(a, x)
+            partial_x = exp_(-x + (a - 1) * log_(x) - lgamma(a))
+            r1 = neg_(reshape_(sum_grad_reduce_axis(partial_a * dout, ra), sa))
+            r2 = neg_(reshape_(sum_grad_reduce_axis(partial_x * dout, rx), sx))
+            return r1, r2
         ra, rx = broadcast_gradient_args(sa, sx)
         partial_a = igammagrada(a, x)
         partial_x = exp_(-x + (a - 1) * log_(x) - lgamma(a))
@@ -1022,6 +1542,63 @@ def get_bprop_igammac(self):
     return bprop
 
 
+@bprop_getters.register(Lgamma)
+def get_bprop_lgamma(self):
+    """Grad definition for `Lgamma` operation."""
+    digamma = Digamma()
+
+    def bprop(x, out, dout):
+        if x.dtype in (mstype.float16,):
+            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,)
+
+    return bprop
+
+
+@bprop_getters.register(Digamma)
+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(a, x)
+            dx = F.cast(dx, mstype.float16)
+        else:
+            dx = dout * polygamma(a, x)
+        return (dx,)
+
+    return bprop
+
+
+@bprop_getters.register(Polygamma)
+def get_bprop_polygamma(self):
+    """Grad definition for `Polygamma` operation."""
+    polygamma = Polygamma()
+
+    def bprop(a, x, out, dout):
+        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 zeros_like(a), dx
+
+    return bprop
+
+
 @bprop_getters.register(TridiagonalMatMul)
 def get_bprop_tridiagonal_matmul(self):
     """Grad definition for 'TridiagonalMatMul' operation"""
@@ -1070,7 +1647,7 @@ def get_bprop_tridiagonal_matmul(self):
         maindiag_grad = reduce_sum(rhs_conj * grad, -1)
         subdiag_grad = reduce_sum(_rightshift(rhs_conj) * grad, -1)
         rhs_grad = _rightshift(superdiag_conj * grad) + maindiag_conj * grad + \
-            _leftshift(subdiag_conj * grad)
+                   _leftshift(subdiag_conj * grad)
         superdiag_grad = expand_dims(superdiag_grad, -2)
         maindiag_grad = expand_dims(maindiag_grad, -2)
         subdiag_grad = expand_dims(subdiag_grad, -2)
@@ -1101,13 +1678,18 @@ def get_bprop_cholesky_solve(self):
 
     def bprop(x1, x2, out, dout):
         flag = 0
+        shape_x1 = shape_op(x1)
+        if F.is_sequence_shape_unknown(shape_x1):
+            len_x1 = dyn_rank(x1)
+        else:
+            len_x1 = len(shape_x1)
         if dout.dtype == mstype.float64:
             flag = 1
             x2 = F.cast(x2, mstype.float32)
             out = F.cast(out, mstype.float32)
             dout = F.cast(dout, mstype.float32)
         dx1 = cholesky_solve(dout, x2)
-        if len(shape_op(x2)) == 2:
+        if len_x1 == 2:
             common_term = matmul_op(dx1, transpose(out, (1, 0)))
             common_term = common_term + transpose(common_term, (1, 0))
             if upper is True:
@@ -1115,8 +1697,11 @@ def get_bprop_cholesky_solve(self):
             else:
                 dx2 = neg_op(matmul_op(common_term, x2))
         else:
-            common_term = batchmatmul_op(dx1, transpose(out, (0, 2, 1)))
-            common_term = common_term + transpose(common_term, (0, 2, 1))
+            x2_dim_size = len(shape_op(x2))
+            x2_dim_order = list(range(x2_dim_size))
+            target_order = x2_dim_order[:-2] + x2_dim_order[-2:][::-1]
+            common_term = batchmatmul_op(dx1, transpose(out, tuple(target_order)))
+            common_term = common_term + transpose(common_term, tuple(target_order))
             if upper is True:
                 dx2 = neg_op(batchmatmul_op(x2, common_term))
             else:
@@ -1133,6 +1718,7 @@ def get_bprop_cholesky_solve(self):
 def get_bprop_nextafter(self):
     """Grad definition for 'NextAfter' operation"""
     shape = P.Shape()
+    dyn_shape = P.TensorShape()
     ones = P.Ones()
     zeros = P.Zeros()
     dtype = P.DType()
@@ -1151,11 +1737,355 @@ def get_bprop_nextafter(self):
             dout = cast(dout, mstype.float32)
 
         s_x1 = shape(x1)
+        partial_x1 = ()
+        if F.is_sequence_value_unknown(s_x1):
+            s_x1 = dyn_shape(x1)
+            partial_x1 = dyn_ones(s_x1, dtype(x1))
+        else:
+            partial_x1 = ones(s_x1, dtype(x1))
+
         s_x2 = shape(x2)
-        partial_x1 = ones(s_x1, dtype(x1))
-        partial_x2 = zeros(s_x2, dtype(x2))
+        partial_x2 = ()
+        if F.is_sequence_value_unknown(s_x2):
+            s_x2 = dyn_shape(x2)
+            partial_x2 = dyn_fill(dtype(x2), s_x2, 0)
+        else:
+            partial_x2 = zeros(s_x2, dtype(x2))
+
         dx1 = reshape(partial_x1 * dout, s_x1)
         dx2 = reshape(partial_x2 * dout, s_x2)
         return cast(dx1, dtype(dout)), cast(dx2, dtype(dout))
 
     return bprop
+
+
+@bprop_getters.register(Diagonal)
+def get_bprop_diagonal(self):
+    """Grad definition for 'Diagonal' operation"""
+    offset = self.offset
+    dim1 = self.dim1
+    dim2 = self.dim2
+    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
+        else:
+            dim1_ = dim1
+        if dim2 < 0:
+            dim2_ = dim2 + x_dim
+        else:
+            dim2_ = dim2
+        if size_op(out):
+            batch_dim = out.shape[:-1]
+            diag_plane = (x_shape[dim1_], x_shape[dim2_])
+            dx_trans_shape = batch_dim + diag_plane
+            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
+            perm = ()
+            for i in range(x_dim):
+                if i == dim1_:
+                    perm = perm + (x_dim - 2,)
+                elif i == dim2_:
+                    perm = perm + (x_dim - 1,)
+                else:
+                    perm = perm + (dim,)
+                    dim = dim + 1
+            dx = transpose_op(dx, perm)
+        else:
+            dx = zeros_like(x)
+        return (dx,)
+
+    return bprop
+
+
+@bprop_getters.register(Cholesky)
+def get_bprop_cholesky(self):
+    """Grad definition for `Cholesky` operation."""
+    upper = self.upper
+    choleskygrad = G.CholeskyGrad()
+
+    def bprop(x, out, dout):
+        out = cholesky_transpose(out) if upper else out
+        dout = cholesky_transpose(dout) if upper else dout
+        dx = choleskygrad(out, dout)
+        return (dx,)
+
+    return bprop
+
+
+@bprop_getters.register(InplaceIndexAdd)
+def get_bprop_inplace_index_add(self):
+    """Generate bprop for InplaceIndexAdd"""
+    gather = P.Gather()
+    _axis = self.axis
+
+    def bprop(var, indices, updates, out, dout):
+        return dout, zeros_like(indices), gather(dout, indices, _axis)
+
+    return bprop
+
+
+@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"""
+    norm_shape_product = 1
+    for i in norm_shape[-rank:]:
+        norm_shape_product *= i
+    return norm_shape_product
+
+
+@_primexpr
+def _fft_with_size_back_norm(norm_shape, norm, inverse, rank):
+    """generate reverse term for fft_with_size"""
+    if inverse is False:
+        if norm == "forward":
+            norm_ = 1 / _fft_rank_offset(norm_shape, rank)
+        if norm == "backward":
+            norm_ = 1 * _fft_rank_offset(norm_shape, rank)
+        if norm == "ortho":
+            norm_ = 1
+    if inverse is True:
+        if norm == "forward":
+            norm_ = 1 * _fft_rank_offset(norm_shape, rank)
+        if norm == "backward":
+            norm_ = 1 / _fft_rank_offset(norm_shape, rank)
+        if norm == "ortho":
+            norm_ = 1
+    return norm_
+
+
+@_primexpr
+def _rfft_norm(norm_shape, norm, rank):
+    """generate norm for rfft"""
+    norm_ = 1.0
+    if norm == "forward":
+        norm_ = 1 / _fft_rank_offset(norm_shape, rank)
+    if norm == "backward":
+        norm_ = 1
+    if norm == "ortho":
+        norm_ = 1 / np.sqrt(_fft_rank_offset(norm_shape, rank))
+    return norm_
+
+
+@_primexpr
+def _get_last_dim_slice_shape(tensor_shape, index):
+    """generate shape for slice last tensor"""
+    from_shape = [0 for x in tensor_shape]
+    if index < 0:
+        from_shape[-1] = tensor_shape[-1] + index
+    else:
+        from_shape[-1] = index
+    to_shape = list(tensor_shape)
+    to_shape[-1] = 1
+    return tuple(from_shape), tuple(to_shape)
+
+
+@_primexpr
+def _rfft_reshape(shape_a, shape_b):
+    """generate rfft shape for reshape"""
+    new_shape = list(shape_b)
+    for i in range(len(shape_a) - 2):
+        new_shape.insert(i, 1)
+    return tuple(new_shape)
+
+
+@_primexpr
+def _rfft_tile_reshape(shape_a):
+    """generate rfft shape for tile"""
+    reshape_a = list(shape_a)
+    reshape_a[-2] = 1
+    reshape_a[-1] = 1
+    return tuple(reshape_a)
+
+
+@_primexpr
+def _rfft_last_term_shape(shape_a, shape_b):
+    """generate rfft shape for last term"""
+    new_shape = list(shape_b)
+    for i in range(len(shape_a) - 1):
+        new_shape.insert(i, 1)
+    return tuple(new_shape)
+
+
+@_primexpr
+def _batch_matmul_shape_increase(shape_before):
+    """increase tensor shape for batch_matmul"""
+    return (1, *shape_before)
+
+
+@_primexpr
+def _batch_matmul_shape_decrease(matrix_shape):
+    """decrease tensor shape after batch_matmul"""
+    shape_tmp = list(matrix_shape)
+    shape_tmp[-1] = 1
+    return tuple(shape_tmp)
+
+
+@bprop_getters.register(FFTWithSize)
+def get_bprop_fft_with_size(self):
+    """Grad definition for `FFTWithSize` operation."""
+    signal_ndim = self.signal_ndim
+    inverse = self.inverse
+    real = self.real
+    norm = self.norm
+    onesided = self.onesided
+    fft_fn = FFTWithSize(signal_ndim=signal_ndim,
+                         inverse=False,
+                         real=False,
+                         norm=norm)
+    ifft_fn = FFTWithSize(signal_ndim=signal_ndim,
+                          inverse=True,
+                          real=False,
+                          norm=norm)
+    rfft_fn = FFTWithSize(signal_ndim=signal_ndim,
+                          inverse=False,
+                          real=True,
+                          norm=norm,
+                          onesided=onesided)
+    irfft_fn = FFTWithSize(signal_ndim=signal_ndim,
+                           inverse=True,
+                           real=True,
+                           norm=norm,
+                           onesided=onesided)
+
+    complex_op = P.Complex()
+    shape_op = P.Shape()
+    to_tensor_op = P.ScalarToTensor()
+    type_op = P.DType()
+    concat_op = P.Concat()
+    ones_op = P.Ones()
+    zeros_op = P.Zeros()
+    real_op = P.Real()
+    imag_op = P.Imag()
+    slice_op = P.Slice()
+    tile_op = P.Tile()
+    expand_dims = P.ExpandDims()
+    transpose_op = P.Transpose()
+    exp_op = P.Exp()
+    reshape_op = P.Reshape()
+    conj_op = P.Conj()
+    batch_matmul_op = P.BatchMatMul()
+
+    def bprop(x, out, dout):
+        dx = 0
+        input_type = type_op(x)
+        output_type = type_op(out)
+        input_shape = shape_op(x)
+        offset_shape = shape_op(x)
+        dout_shape = shape_op(dout)
+        offset_size = to_tensor_op(_fft_with_size_back_norm(offset_shape, norm, inverse, signal_ndim), output_type)
+        if real is False:
+            if inverse is False:
+                dx = ifft_fn(dout) * offset_size
+            else:
+                dx = fft_fn(dout) * offset_size
+        else:
+            irfft_ = FFTWithSize(signal_ndim=1, inverse=True, real=True, norm="backward", onesided=onesided,
+                                 signal_sizes=offset_shape[-1:])
+            irfft2d_ = FFTWithSize(signal_ndim=2, inverse=True, real=True, norm="backward", onesided=onesided,
+                                   signal_sizes=offset_shape[-2:])
+            irfft3d_ = FFTWithSize(signal_ndim=3, inverse=True, real=True, norm="backward", onesided=onesided,
+                                   signal_sizes=offset_shape[-3:])
+            if inverse is False:
+                if onesided is True:
+                    terms = 0
+                    is_even = to_tensor_op(1 - (input_shape[-1] % 2), input_type)
+                    dout_first_from, dout_first_to = _get_last_dim_slice_shape(dout_shape, 0)
+                    dout_first = slice_op(dout, dout_first_from, dout_first_to)
+                    rfft_offset_size = to_tensor_op(_fft_rank_offset(input_shape, signal_ndim), input_type)
+                    rfft_norm_offset = to_tensor_op(_rfft_norm(input_shape, norm, signal_ndim), input_type)
+                    dout_last_from, dout_last_to = _get_last_dim_slice_shape(dout_shape, -1)
+                    dout_last = slice_op(dout, dout_last_from, dout_last_to)
+                    if signal_ndim == 1:
+                        dx = irfft_(dout)
+                        vec_mask = complex_op(1 - 2 * (mnp.arange(0, input_shape[-1], 1, input_type) % 2),
+                                              zeros_op(input_shape[-1], input_type))
+                        terms = real_op(dout_first) + is_even * real_op(dout_last * vec_mask)
+                    elif signal_ndim == 2:
+                        dx = irfft2d_(dout)
+                        arange_inner = mnp.arange(0, input_shape[-2], 1, input_type)
+                        matrix_a = tile_op(expand_dims(arange_inner, 0), (input_shape[-2], 1))
+                        matrix_b = transpose_op(matrix_a, (1, 0))
+                        matrix_mul = matrix_a * matrix_b
+                        imag_offset = complex_op(to_tensor_op(0, input_type), to_tensor_op(-2, input_type))
+                        pi_tensor = to_tensor_op(mnp.pi, output_type)
+                        matrix_mul_complex = complex_op(matrix_mul, zeros_op(shape_op(matrix_mul), input_type))
+                        matrix_base_mask = exp_op(imag_offset * pi_tensor * matrix_mul_complex /
+                                                  to_tensor_op(input_shape[-2], output_type))
+                        expanded_matrix_mask = reshape_op(matrix_base_mask, _rfft_reshape(shape_op(dout_first),
+                                                                                          shape_op(matrix_base_mask)))
+                        tile_matrix_mask = complex_op(tile_op(real_op(expanded_matrix_mask), _rfft_tile_reshape(
+                            shape_op(dout_first))), tile_op(imag_op(expanded_matrix_mask),
+                                                            _rfft_tile_reshape(shape_op(dout_first))))
+                        tile_matrix_mask_shape = shape_op(tile_matrix_mask)
+                        dout_first_term = reshape_op(batch_matmul_op(reshape_op(tile_matrix_mask,
+                                                                                _batch_matmul_shape_increase(
+                                                                                    tile_matrix_mask_shape)),
+                                                                     reshape_op(conj_op(
+                                                                         dout_first), _batch_matmul_shape_increase(
+                                                                             shape_op(dout_first)))),
+                                                     _batch_matmul_shape_decrease(tile_matrix_mask_shape))
+                        dout_last_term = reshape_op(batch_matmul_op(reshape_op(tile_matrix_mask,
+                                                                               _batch_matmul_shape_increase(
+                                                                                   tile_matrix_mask_shape)),
+                                                                    reshape_op(conj_op(dout_last),
+                                                                               _batch_matmul_shape_increase(
+                                                                                   shape_op(dout_last)))),
+                                                    _batch_matmul_shape_decrease(
+                                                        tile_matrix_mask_shape))
+                        vec_mask = complex_op(1 - 2 * (mnp.arange(0, input_shape[-1], 1, input_type) % 2), zeros_op(
+                            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]])),
+                                                    tile_op(imag_op(dout_last_term), _rfft_last_term_shape(
+                                                        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:
+                        dx = irfft3d_(dout) * real_op(offset_size)
+                    dx = to_tensor_op(0.5, input_type) * (dx * rfft_offset_size + terms) * rfft_norm_offset
+                else:
+                    dx = irfft_fn(dout) * real_op(offset_size)
+            else:
+                dx = rfft_fn(dout)
+                if onesided is True:
+                    if signal_ndim != 3:
+                        is_odd = dout_shape[-1] % 2
+                        last_shape = offset_shape[-1]
+                        mask = concat_op((ones_op(1, output_type), 2.0 * ones_op(
+                            (last_shape - 2 + is_odd,), output_type), ones_op((1 - is_odd,), output_type)))
+                        dx = dx * complex_op(mask, zeros_op(shape_op(mask), output_type))
+                        irfft_offset_size = to_tensor_op(
+                            _fft_with_size_back_norm(shape_op(dout), norm, inverse, signal_ndim),
+                            output_type)
+                        dx = dx * complex_op(irfft_offset_size, zeros_op(1, output_type))
+                    else:
+                        dx = dx * complex_op(offset_size, zeros_op(1, output_type))
+                else:
+                    dx = dx * complex_op(offset_size, zeros_op(1, output_type))
+        return (dx,)
+
+    return bprop