mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

mindspore/ops/_grad_experimental/grad_comm_ops.py

@@ -22,7 +22,8 @@ from mindspore.ops import functional as F
 from mindspore.communication import get_rank, get_group_size
 from mindspore.parallel._utils import _get_enable_parallel_optimizer, _get_grad_accumulation_shard
 from mindspore.ops import operations as P
-from mindspore.ops.operations._inner_ops import Send, Receive, issubclass_
+from mindspore.ops import Send, Receive
+from mindspore.ops.operations._inner_ops import issubclass_
 from mindspore.common.sparse_tensor import RowTensorInner
 from mindspore.ops.composite.multitype_ops.zeros_like_impl import zeros_like
 from mindspore.ops.operations.comm_ops import (AllGather, _MiniStepAllGather, _HostAllGather, AllReduce,
@@ -30,7 +31,7 @@ from mindspore.ops.operations.comm_ops import (AllGather, _MiniStepAllGather, _H
                                                _GetTensorSlice, _MirrorOperator, _MirrorMiniStepOperator, ReduceOp,
                                                ReduceScatter, _HostReduceScatter, _VirtualDiv, _VirtualAdd, _AllSwap,
                                                _VirtualAssignAdd, _VirtualAccuGrad, _MirrorMicroStepOperator,
-                                               _MicroStepAllGather)
+                                               _MicroStepAllGather, Reduce, CollectiveGather, CollectiveScatter)
 from mindspore.ops._grad_experimental.grad_base import bprop_getters
 from mindspore.ops.operations import _grad_ops as G
 
@@ -41,7 +42,7 @@ def get_bprop_all_reduce(self):
 
     all_reduce_grad = AllReduce(ReduceOp.SUM, self.group)
     all_gather = AllGather(group=self.group)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad" + self.instance_name
         all_reduce_grad.set_prim_instance_name(instance_name)
     equal = P.Equal()
@@ -108,6 +109,7 @@ def get_bprop_receive(self):
     """Generate bprop for Receive."""
     tag = self.get_attr_dict()["sr_tag"]
     receive_grad = Send(tag, self.rank, self.group_back)
+    receive_grad.add_prim_attr("shape", self.shape)
     depend = P.Depend()
     cast = P.Cast()
     out_tensor = Tensor(0.0, mstype.float16)
@@ -209,21 +211,17 @@ def get_bprop_mirror_micro_step_operator(self):
     def bprop(x, z, out, dout):
         real_grad = z
         assign_out = dout
-        if mean_flag:
-            if issubclass_(F.typeof(dout), mstype.tensor_type):
-                z = F.depend(z, dout)
-                real_grad = all_reduce(z)
-                real_grad = F.tensor_mul(real_grad, scale)
-                if opt_shard:
-                    return (real_grad, cast(out_tensor, dtype(z)))
-                return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign(z, real_grad))
-        else:
-            if issubclass_(F.typeof(dout), mstype.tensor_type):
-                z = F.depend(z, dout)
+        if issubclass_(F.typeof(dout), mstype.tensor_type):
+            z = F.depend(z, dout)
+            if dev_num > 1:
                 real_grad = all_reduce(z)
-                if opt_shard:
-                    return (real_grad, cast(out_tensor, dtype(z)))
-                return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign(z, real_grad))
+                if mean_flag:
+                    real_grad = F.tensor_mul(real_grad, scale)
+            else:
+                real_grad = z
+            if opt_shard:
+                return (real_grad, cast(out_tensor, dtype(z)))
+            return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign(z, real_grad))
         return F.depend((cast(out_tensor, dtype(x)), cast(out_tensor, dtype(z))), assign_out)
 
     return bprop
@@ -243,11 +241,13 @@ def get_bprop_broad_cast(self):
 def get_bprop_all_gather(self):
     """Generate bprop for AllGather"""
     fusion = self.get_attr_dict()["fusion"]
+    self.group = self.get_attr_dict()["group"]
     reduce_scatter = ReduceScatter(ReduceOp.SUM, self.group).add_prim_attr("fusion", fusion)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad_" + self.instance_name
         reduce_scatter.set_prim_instance_name(instance_name)
     mean_flag = self.get_attr_dict()["mean_flag"]
+    self.rank_size = self.get_attr_dict()["rank_size"]
     if self.rank_size == 0:
         raise ValueError(f"The 'rank_size' can not be zero, but got {self.rank_size}.")
     scale = 1.0 / self.rank_size
@@ -272,7 +272,7 @@ def get_bprop_mini_step_all_gather(self):
     scale = 1 / self.rank_size
     all_reduce = AllReduce(ReduceOp.SUM, self.group).add_prim_attr("fusion", fusion)
     assign_add = P.AssignAdd()
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad_" + self.instance_name
         all_reduce.set_prim_instance_name(instance_name)
     rank = get_rank(self.group)
@@ -314,15 +314,12 @@ def get_bprop_micro_step_all_gather(self):
         do_mirror = self.get_attr_dict()["do_mirror"]
     if do_mirror:
         scale = 1.0 / self.rank_size
-        all_reduce = AllReduce(ReduceOp.SUM, self.group).add_prim_attr("fusion", fusion)
+        reduce_scatter = ReduceScatter(ReduceOp.SUM, self.group).add_prim_attr("fusion", fusion)
         if "segment" in self.get_attr_dict():
-            all_reduce.add_prim_attr("segment", self.get_attr_dict()["segment"])
-        rank = get_rank(self.group)
-        dev_num = get_group_size(self.group)
-        split = P.Split(output_num=dev_num)
+            reduce_scatter.add_prim_attr("segment", self.get_attr_dict()["segment"])
         if self.instance_name:
             instance_name = "grad_" + self.instance_name
-            all_reduce.set_prim_instance_name(instance_name)
+            reduce_scatter.set_prim_instance_name(instance_name)
     cast = P.Cast()
     dtype = P.DType()
     out_tensor = Tensor(1.0, mstype.float16)
@@ -332,16 +329,14 @@ def get_bprop_micro_step_all_gather(self):
         if with_mirror_operator:
             if not do_mirror:
                 return (dout, cast(out_tensor, dtype(z)))
-            real_grad = all_reduce(dout)
-            real_grad = split(real_grad)[rank]
+            real_grad = reduce_scatter(dout)
             if mean_flag:
                 real_grad = F.tensor_mul(real_grad, scale)
             return (real_grad, cast(out_tensor, dtype(z)))
         z = F.depend(z, dout)
         if not do_mirror:
             return (z, cast(out_tensor, dtype(z)))
-        real_grad = all_reduce(z)
-        real_grad = split(real_grad)[rank]
+        real_grad = reduce_scatter(z)
         if mean_flag:
             real_grad = F.tensor_mul(real_grad, scale)
         return (real_grad, cast(out_tensor, dtype(z)))
@@ -353,7 +348,7 @@ def get_bprop_micro_step_all_gather(self):
 def get_bprop_host_all_gather(self):
     """Generate bprop for _HostAllGather"""
     host_all_gather_grad = _HostReduceScatter(ReduceOp.SUM, self.group)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad" + self.instance_name
         host_all_gather_grad.set_prim_instance_name(instance_name)
 
@@ -368,7 +363,7 @@ def get_bprop_host_all_gather(self):
 def get_bprop_reduce_scatter(self):
     """Generate bprop for ReduceScatter"""
     reduce_scatter_grad = AllGather(self.group)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad" + self.instance_name
         reduce_scatter_grad.set_prim_instance_name(instance_name)
 
@@ -382,11 +377,71 @@ def get_bprop_reduce_scatter(self):
     return bprop
 
 
+@bprop_getters.register(Reduce)
+def get_bprop_reduce(self):
+    """Generate bprop for Reduce"""
+    dest_rank = self.get_attr_dict()["dest_rank"]
+    group = self.get_attr_dict()["group"]
+    reduce_grad = Broadcast(dest_rank, group)
+    if hasattr(self, "instance_name") and self.instance_name:
+        instance_name = "grad" + self.instance_name
+        reduce_grad.set_prim_instance_name(instance_name)
+
+    def bprop(x, out, dout):
+        dx = reduce_grad((dout,))
+        return (dx[0],)
+
+    return bprop
+
+
+@bprop_getters.register(CollectiveGather)
+def get_bprop_collective_gather(self):
+    """Generate bprop for CollectiveGather"""
+    group = self.get_attr_dict()["group"]
+    dest_rank = self.get_attr_dict()["dest_rank"]
+    collective_gather_grad = Broadcast(dest_rank, group)
+    rank = get_rank(group)
+    dev_num = self.rank_size
+    split = P.Split(output_num=dev_num)
+    if hasattr(self, "instance_name") and self.instance_name:
+        instance_name = "grad" + self.instance_name
+        collective_gather_grad.set_prim_instance_name(instance_name)
+
+    def bprop(x, out, dout):
+        grad = collective_gather_grad((dout,))
+        dx = split(grad[0])[rank]
+        return (dx,)
+
+    return bprop
+
+
+@bprop_getters.register(CollectiveScatter)
+def get_bprop_collective_scatter(self):
+    """Generate bprop for CollectiveScatter"""
+    group = self.get_attr_dict()["group"]
+    dest_rank = self.get_attr_dict()["src_rank"]
+    rank = get_rank(group)
+    collective_scatter_grad = CollectiveGather(dest_rank, group)
+    if hasattr(self, "instance_name") and self.instance_name:
+        instance_name = "grad" + self.instance_name
+        collective_scatter_grad.set_prim_instance_name(instance_name)
+
+    def bprop(x, out, dout):
+        dx_out = collective_scatter_grad(dout)
+        if rank == dest_rank:
+            dx = dx_out
+        else:
+            dx = F.depend(F.zeros_like(x), dx_out)
+        return (dx,)
+
+    return bprop
+
+
 @bprop_getters.register(_AllSwap)
 def get_bprop_allswap(self):
     """Generate bprop for _AllSwap."""
     all_swap_grad = _AllSwap(self.group)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad" + self.instance_name
         all_swap_grad.set_prim_instance_name(instance_name)
 
@@ -401,7 +456,7 @@ def get_bprop_allswap(self):
 def get_bprop_host_reduce_scatter(self):
     """Generate bprop for _HostReduceScatter"""
     host_reduce_scatter_grad = _HostAllGather(self.group)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad" + self.instance_name
         host_reduce_scatter_grad.set_prim_instance_name(instance_name)
 
@@ -436,7 +491,7 @@ def get_bprop_neighborexchange(self):
 def get_bprop_all_to_all(self):
     """Generate bprop for AlltoAll."""
     all_to_all_grad = AlltoAll(self.split_count, self.concat_dim, self.split_dim, self.group)
-    if self.instance_name:
+    if hasattr(self, "instance_name") and self.instance_name:
         instance_name = "grad" + self.instance_name
         all_to_all_grad.set_prim_instance_name(instance_name)
 

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -22,13 +22,8 @@ import mindspore.ops as ops
 from mindspore.ops import functional as F
 from mindspore.ops import operations as P
 from mindspore import Tensor
-from mindspore.ops.operations.math_ops import Real, Imag, Complex, Angle
-from mindspore.ops.operations.math_ops import Polar, SilentCheck
+from mindspore.ops.operations.math_ops import SilentCheck
 from mindspore.ops.operations._inner_ops import _MirrorSilentCheck
-from mindspore.ops.operations import _grad_ops as G
-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 CumulativeLogsumexp
 from mindspore.ops.operations.math_ops import MatrixSolve
 from mindspore.ops.operations.math_ops import MatrixSolveLs
@@ -37,11 +32,8 @@ from mindspore.ops.operations.math_ops import NanToNum
 from mindspore.ops.operations.math_ops import FFTWithSize
 from mindspore.ops.operations.math_ops import Cholesky
 from mindspore.ops.operations.math_ops import CholeskySolve
-from mindspore.ops.operations.math_ops import InplaceIndexAdd
-from mindspore.ops.operations.math_ops import TridiagonalSolve
-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.array_ops import Transpose
 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 _primexpr
@@ -137,7 +129,7 @@ def get_bprop_matrix_triangular_solve(self):
 
     def bprop(matrix, rhs, out, dout):
         grad_rhs = matrix_triangular_solve_op(matrix, dout)
-        if matrix.dtype in (mstype.complex64, mstype.complex128):
+        if matrix.dtype == mstype.complex64 or matrix.dtype == mstype.complex128:
             grad_rhs_temp = _adjoint(grad_rhs)
             out_temp = _adjoint(out)
         else:
@@ -158,14 +150,14 @@ def get_bprop_matrix_triangular_solve(self):
                 grad_matrix = mat_mul_op(grad_rhs, out_temp)
                 grad_matrix = neg_op(grad_matrix)
         if lower_a:
-            if grad_matrix.dtype in (mstype.complex64, mstype.complex128):
+            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 in (mstype.complex64, mstype.complex128):
+            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)
@@ -221,7 +213,7 @@ def get_bprop_matrix_solve(self):
 @_primexpr
 def _generate_perm_matrix_solve_ls(x_dim):
     perm = tuple(range(x_dim - 2))
-    perm = perm + (x_dim - 1, x_dim - 2)
+    perm = perm + (x_dim-1, x_dim-2)
     return perm
 
 
@@ -285,7 +277,7 @@ def get_bprop_matrix_solve_ls(self):
             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
+            # CholeskySolve not support complex dtype and just support 2D or 3D matrices for now
             z = cholesky_solve(dout, chol)
 
             matrix_dim = rank(matrix)
@@ -380,131 +372,6 @@ def get_bprop_nan_to_num(self):
     return bprop
 
 
-@bprop_getters.register(Angle)
-def get_bprop_angle(self):
-    """Grad definition for `Angle` operation."""
-    real_op = Real()
-    imag_op = Imag()
-    reciprocal_op = P.Reciprocal()
-    complex_op = Complex()
-    neg_op = P.Neg()
-
-    def bprop(x, out, dout):
-        re = real_op(x)
-        im = imag_op(x)
-        re = complex_op(im, re)
-        z = reciprocal_op(re)
-        zero = zeros_like(dout)
-        complex_dout = complex_op(dout, zero)
-        return (neg_op(complex_dout * z),)
-
-    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(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(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(CholeskySolve)
 def get_bprop_cholesky_solve(self):
     """Grad definition for 'CholeskySolve' operation"""
@@ -549,94 +416,6 @@ def get_bprop_cholesky_solve(self):
     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"""
@@ -646,24 +425,37 @@ def _fft_rank_offset(norm_shape, rank):
     return norm_shape_product
 
 
+@_primexpr
+def _norm_enum_to_string(norm_mode):
+    """convert norm_mode enum to string ."""
+    norm_mode_str = ""
+    if norm_mode == 0:
+        norm_mode_str = "backward"
+    elif norm_mode == 1:
+        norm_mode_str = "forward"
+    elif norm_mode == 2:
+        norm_mode_str = "ortho"
+    return norm_mode_str
+
+
 @_primexpr
 def _fft_with_size_back_norm(norm_shape, norm, inverse, rank):
     """generate reverse term for fft_with_size"""
-    norm_ = None
+    norm_ = 1
     if inverse is False:
         if norm == "forward":
-            norm_ = 1.0 / _fft_rank_offset(norm_shape, rank)
-        elif norm == "backward":
-            norm_ = 1.0 * _fft_rank_offset(norm_shape, rank)
-        elif norm == "ortho":
-            norm_ = 1.0
-    else:
+            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.0 * _fft_rank_offset(norm_shape, rank)
-        elif norm == "backward":
-            norm_ = 1.0 / _fft_rank_offset(norm_shape, rank)
-        elif norm == "ortho":
-            norm_ = 1.0
+            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_
 
 
@@ -673,9 +465,9 @@ def _rfft_norm(norm_shape, norm, rank):
     norm_ = 1.0
     if norm == "forward":
         norm_ = 1 / _fft_rank_offset(norm_shape, rank)
-    elif norm == "backward":
-        norm_ = 1.0
-    elif norm == "ortho":
+    if norm == "backward":
+        norm_ = 1
+    if norm == "ortho":
         norm_ = 1 / np.sqrt(_fft_rank_offset(norm_shape, rank))
     return norm_
 
@@ -737,30 +529,6 @@ def _batch_matmul_shape_decrease(matrix_shape):
 @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()
     to_tensor_op = P.ScalarToTensor()
     type_op = P.DType()
@@ -779,7 +547,31 @@ def get_bprop_fft_with_size(self):
     conj_op = P.Conj()
     batch_matmul_op = P.BatchMatMul()
 
-    def bprop(x, out, dout):
+    def reverse_branch(dx, onesided, dout_shape, offset_shape,
+                       output_type, dout, norm, inverse, signal_ndim, offset_size):
+        if onesided is True:
+            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))
+        return dx
+
+    def bprop(x, signal_ndim, inverse, real, norm_enum, onesided, signal_sizes, out, dout):
+        norm = _norm_enum_to_string(norm_enum)
+        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)
         dx = 0
         input_type = type_op(x)
         output_type = type_op(out)
@@ -866,18 +658,8 @@ def get_bprop_fft_with_size(self):
                     dx = irfft_fn(dout) * real_op(offset_size)
             else:
                 dx = rfft_fn(dout)
-                if onesided is True:
-                    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))
+                dx = reverse_branch(dx, onesided, dout_shape, offset_shape,
+                                    output_type, dout, norm, inverse, signal_ndim, offset_size)
         return (dx,)
 
     return bprop
@@ -1030,10 +812,13 @@ def get_bprop_mirror_silent_check(self):
     out_tensor = Tensor([0.0], mstype.float32)
 
     def bporp(x, pre_val, min_val, max_val, n_step, loss_scale, out, dout):
+        if dout.dtype == mstype.float16:
+            return (dout, out_tensor, out_tensor, out_tensor, out_tensor, out_tensor)
         if loss_scale is not None:
-            dout = dout / loss_scale
-        grad = ops.norm(dout)
-        dx, _, _, _, _ = silent_check(grad, dout, pre_val, min_val, max_val, n_step)
+            gnorm = ops.norm(dout / loss_scale)
+        else:
+            gnorm = ops.norm(dout)
+        dx, _, _, _, _ = silent_check(gnorm, dout, pre_val, min_val, max_val, n_step)
         return (dx, out_tensor, out_tensor, out_tensor, out_tensor, out_tensor)
 
     return bporp