mindspore 2.2.10__cp38-none-any.whl → 2.2.14__cp38-none-any.whl

mindspore/nn/layer/flash_attention.py

@@ -21,9 +21,7 @@ import mindspore.common.dtype as mstype
 from mindspore.common.tensor import Tensor
 from mindspore import ops
 from mindspore.nn.cell import Cell
-from mindspore.ops._op_impl._custom_op.flash_attention.flash_attention_impl import get_flash_attention
 from mindspore.ops.operations.nn_ops import FlashAttentionScore
-from mindspore._c_expression import MSContext
 
 __all__ = ['FlashAttention']
 
@@ -46,18 +44,17 @@ class FlashAttention(Cell):
             Default 65536.
         next_block_num(int): A integer to define the number of blocks to look behind for local block sparse attention.
             Default 65536.
-        tiling_stgy_name(str): A str to define tiling strategy of flash attention.
         dp(int): data parallel.
             Default 1.
         mp(int): model parallel.
             Default 1.
-        high_precision(bool): This mode has higher precision but some performance loss.
+        high_precision(bool): This mode has higher precision but some performance loss. Only take effect on Ascend910A.
             Default False.
         have_attention_mask_batch(bool): indicates whether attention_mask contains the batch dimension.
             Default True
         alibi(bool): This parameter indicates whether the flashattention supports the Alibi.
             Default: False
-        use_mqa(bool): Using MHA if True, only take effect under 910B. Default: False.
+        use_mqa(bool): Using MQA if True, only take effect under 910B. Default: False.
 
 
     Inputs:
@@ -98,7 +95,6 @@ class FlashAttention(Cell):
                  dropout_rate=0.0,
                  prev_block_num=65536,
                  next_block_num=65536,
-                 tiling_stgy_name="sparse",
                  dp=1,
                  mp=1,
                  high_precision=False,
@@ -112,52 +108,36 @@ class FlashAttention(Cell):
         if scaling_constant == 0:
             raise ValueError("the scaling constant must not be 0.")
         self.dropout_rate = dropout_rate
-        self.is_910A = MSContext.get_instance().get_ascend_soc_version() == "ascend910"
-        if self.is_910A:
-            self.scale_factor = Tensor([1. / math.sqrt(scaling_constant)], dtype=mstype.float16)
-            self.scale_mul = ops.Mul().shard(((dp, mp, 1, 1), (1,)))
-            self.ones = ops.Ones()
-            self.dim_mask = Tensor([1 for _ in range(head_dim)], dtype=mstype.int8)
-            self.have_attention_mask_batch = have_attention_mask_batch
-            self.alibi = alibi
-            self.flash_attention = get_flash_attention(
-                prev_block_num=prev_block_num,
-                next_block_num=next_block_num,
-                tiling_stgy_name=tiling_stgy_name,
-                high_precision=high_precision
-            )
-            self.flash_attention.add_prim_attr("primitive_target", "Ascend")
+        self.alibi = alibi
+        self.have_attention_mask_batch = have_attention_mask_batch
+
+        self.transpose_4d_pre = ops.Transpose().shard(((dp, mp, 1, 1),))
+        self.transpose_4d_post = ops.Transpose().shard(((dp, 1, mp, 1),))
+        self.reshape = ops.Reshape()
+        self.zeros_like = ops.ZerosLike().shard(((dp, mp, 1, 1),))
+        self.zeros = ops.Zeros()
+        self.attn_cast = ops.Cast()
+        if use_mqa:
+            fa_strategies = ((dp, mp, 1, 1),
+                             (dp, 1, 1, 1),
+                             (dp, 1, 1, 1))
+        else:
             fa_strategies = ((dp, mp, 1, 1),
                              (dp, mp, 1, 1),
                              (dp, mp, 1, 1))
-            self.shard(fa_strategies)
-        else:
-            if alibi:
-                raise ValueError(f"When soc_version is not Ascend910A, alibi must be False")
-            self.transpose_4d_pre = ops.Transpose().shard(((dp, mp, 1, 1),))
-            self.transpose_4d_post = ops.Transpose().shard(((dp, 1, mp, 1),))
-            self.reshape = ops.Reshape()
-            self.zeros_like = ops.ZerosLike().shard(((dp, mp, 1, 1),))
-            self.zeros = ops.Zeros()
-            self.attn_cast = ops.Cast()
-            if use_mqa:
-                fa_strategies = ((dp, mp, 1, 1),
-                                 (dp, 1, 1, 1),
-                                 (dp, 1, 1, 1),
-                                 (dp, 1, 1, 1))
-            else:
-                fa_strategies = ((dp, mp, 1, 1),
-                                 (dp, mp, 1, 1),
-                                 (dp, mp, 1, 1),
-                                 (dp, 1, 1, 1))
-            if dropout_rate > 1e-5:
-                fa_strategies += ((dp, mp, 1, 1),)
-            self.flash_attention = FlashAttentionScore(head_num=head_num, pre_tokens=prev_block_num,
-                                                       next_tokens=next_block_num,
-                                                       keep_prob=1 - dropout_rate,
-                                                       scale_value=1. / scaling_constant,
-                                                       inner_precise=0 if high_precision else 1,
-                                                       input_layout="BNSD").shard(fa_strategies)
+        if self.alibi:
+            self.alibi_rescale_mul = ops.Mul().shard(((dp, mp, 1, 1), (1,)))
+            self.alibi_rescale_factor = Tensor([scaling_constant], dtype=mstype.float16)
+            fa_strategies += ((dp, mp, 1, 1),)
+        if dropout_rate > 1e-5:
+            fa_strategies += ((dp, mp, 1, 1),)
+        fa_strategies += ((dp, 1, 1, 1),)
+        self.flash_attention = FlashAttentionScore(head_num=head_num, pre_tokens=prev_block_num,
+                                                   next_tokens=next_block_num,
+                                                   keep_prob=1 - dropout_rate,
+                                                   scale_value=1. / scaling_constant,
+                                                   inner_precise=0,
+                                                   input_layout="BNSD").shard(fa_strategies)
 
         self.dropout_rate = dropout_rate
         if self.dropout_rate > 1e-5:
@@ -175,49 +155,7 @@ class FlashAttention(Cell):
                                   such as MatMul. Default: None.
         :return:
         """
-        if self.is_910A:
-            if in_strategy is None:
-                # default: dp=1, mp=1, construct inputs only contain query, key, value
-                in_strategy = (
-                    (1, 1, 1, 1),
-                    (1, 1, 1, 1),
-                    (1, 1, 1, 1),
-                )
-            self.flash_attention.shard(in_strategy)
-            dp = in_strategy[0][0]
-            mp = in_strategy[0][1]
-            self.flash_attention.add_prim_attr("dev_matrix_shape", [dp, mp, 1, 1])
-            inputs_tensor_map = [
-                [3, 2, 1, 0],
-                [3, 2, 1, 0],
-                [3, 2, 1, 0],
-            ]
-            if self.have_attention_mask_batch:
-                inputs_tensor_map.append([3, 1, 0])
-            else:
-                inputs_tensor_map.append([-1, 1, 0])
-
-            input_empty_args_num = 2
-            # dropout_mask
-            if self.dropout_rate > 1e-5:
-                input_empty_args_num -= 1
-                inputs_tensor_map.append([3, 2, 1, 0])
-
-            if self.alibi:
-                input_empty_args_num -= 1
-                inputs_tensor_map.append([3, 2, 1, 0])
-
-            self.flash_attention.add_prim_attr("inputs_tensor_map", inputs_tensor_map)
-
-            self.flash_attention.add_prim_attr("outputs_tensor_map", [
-                [3, 2, 1, 0],  # O
-                [3, 2, 1],  # L
-                [3, 2, 1]  # M
-            ])
-            self.flash_attention.add_prim_attr("as_loss_divisor", 0)
-            self.flash_attention.add_prim_attr("empty_mirror_ops", input_empty_args_num)
-        else:
-            self.flash_attention.shard(in_strategy)
+        self.flash_attention.shard(in_strategy)
 
     def construct(self, query, key, value, attn_mask=None, alibi_mask=None):
         """FlashAttention forward
@@ -228,49 +166,24 @@ class FlashAttention(Cell):
         :param alibi_mask: [bsz, head_num, 1, seq_len], if not None
         :return: output          [bsz, head_num, seq_len, head_dim]
         """
-        bsz, head_num, seq_len, head_dim = query.shape
-        if self.is_910A:
-            _, k_head_num, k_seq_len, _ = key.shape
-            _, v_head_num, v_seq_len, _ = value.shape
-            if head_num != k_head_num or head_num != v_head_num:
-                raise ValueError(
-                    "the head_num of query, key and value must be the same, "
-                    "If different head_num are used, users need to change themselves to be same by tile.")
-            if seq_len % 16 != 0 or k_seq_len % 16 != 0 or k_seq_len != v_seq_len:
-                raise ValueError(
-                    "query, key, value seq_len must be a multiple of 16, "
-                    "and the seq_len between key and value must be equal.")
-            # 910A -- FlashAttentionPrimtive
-            if head_dim > 304:
-                raise ValueError(
-                    "the head_dim must be less than 304, otherwise the ub would be OOM.")
-            if self.dropout_rate > 1e-5:
-                drop_mask_bits = self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob)
-                tensor_shape = Tensor((bsz, head_num, seq_len, seq_len), mstype.int32)
-                ones = self.fill_v2(tensor_shape, self.tensor_one)
-                ones = self.depend(ones, query)
-                drop_mask = self.do_dropout(ones, drop_mask_bits, self.keep_prob)
-            else:
-                drop_mask = None
-            query = self.scale_mul(query, self.scale_factor)
-            key = self.scale_mul(key, self.scale_factor)
-            attn_mask = self.cast(attn_mask, mstype.float16)
-            output, _, _ = self.flash_attention(query, key, value, attn_mask, drop_mask, alibi_mask)
+        bsz, head_num, seq_len, _ = query.shape
+        # 910B -- FlashAttentionScore
+        if self.dropout_rate > 1e-5:
+            drop_mask_bits = self.reshape(self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob),
+                                          (bsz, head_num, seq_len, seq_len // 8))
         else:
-            # 910B -- FlashAttentionScore
-            if self.dropout_rate > 1e-5:
-                drop_mask_bits = self.reshape(self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob),
-                                              (bsz, head_num, seq_len, seq_len // 8))
-            else:
-                drop_mask_bits = None
-            # (B, S, S) -> (B, 1, S, S)
+            drop_mask_bits = None
+        if self.alibi:
+            alibi_mask = self.alibi_rescale_mul(alibi_mask, self.cast(self.alibi_rescale_factor, alibi_mask.dtype))
+        # (B, S, S) -> (B, 1, S, S)
+        if self.have_attention_mask_batch:
             attn_mask = self.cast(self.reshape(attn_mask, (bsz, 1, seq_len, seq_len)), mstype.uint8)
-            output, _, _ = self.flash_attention(query,
-                                                key,
-                                                value,
-                                                attn_mask,
-                                                drop_mask_bits,
-                                                None,
-                                                None,
-                                                None)
+        _, _, _, output = self.flash_attention(query,
+                                               key,
+                                               value,
+                                               alibi_mask,
+                                               drop_mask_bits,
+                                               None,
+                                               attn_mask,
+                                               None)
         return output

mindspore/nn/loss/loss.py

@@ -1996,7 +1996,7 @@ class FocalLoss(LossBase):
         >>> import mindspore.nn as nn
         >>> logits = ms.Tensor([[0.8, 1.4], [0.5, 0.9], [1.2, 0.9]], ms.float32)
         >>> labels = ms.Tensor([[1], [1], [0]], ms.int32)
-        >>> focalloss = nn.FocalLoss(weight=Tensor([1, 2]), gamma=2.0, reduction='mean')
+        >>> focalloss = nn.FocalLoss(weight=ms.Tensor([1, 2]), gamma=2.0, reduction='mean')
         >>> output = focalloss(logits, labels)
         >>> print(output)
         0.12516622

mindspore/nn/optim/ada_grad.py

@@ -162,7 +162,7 @@ class Adagrad(Optimizer):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-        >>> import mindspore
+        >>> from mindspore import train
         >>> import mindspore.nn as nn
         >>>
         >>> # Define the network structure of LeNet5. Refer to
@@ -185,7 +185,7 @@ class Adagrad(Optimizer):
         >>> # The final parameters order in which the optimizer will be followed is the value of 'order_params'.
         >>>
         >>> loss = nn.SoftmaxCrossEntropyWithLogits()
-        >>> model = ms.Model(net, loss_fn=loss, optimizer=optim)
+        >>> model = train.Model(net, loss_fn=loss, optimizer=optim)
     """
 
     @opt_init_args_register

mindspore/nn/optim/sgd.py

@@ -193,9 +193,9 @@ class SGD(Optimizer):
                             "or 'weight_decay' set in grouped 'params' must be float or int type.")
 
         if hasattr(self, "group_weight_decay") and self.group_weight_decay:
-            self.opt = tuple(P.SGD(dampening, wd, nesterov) for wd in self.group_weight_decay)
+            self.opt = tuple(P.SGD(dampening, 0.0, nesterov) for _ in self.group_weight_decay)
         else:
-            self.opt = tuple([P.SGD(dampening, float(weight_decay), nesterov)] * len(self._parameters))
+            self.opt = tuple([P.SGD(dampening, 0.0, nesterov)] * len(self._parameters))
 
         self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
 
@@ -220,6 +220,7 @@ class SGD(Optimizer):
         params = self._parameters
         accum = self.accum
         stat = self.stat
+        gradients = self.decay_weight(gradients)
         gradients = self.flatten_gradients(gradients)
         gradients = self.gradients_centralization(gradients)
         gradients = self.scale_grad(gradients)

mindspore/nn/wrap/__init__.py

@@ -20,7 +20,8 @@ Use the Wrapper to combine the loss or build the training steps.
 from __future__ import absolute_import
 
 from mindspore.nn.wrap.cell_wrapper import ForwardValueAndGrad, TrainOneStepCell, WithLossCell, WithGradCell, \
-    WithEvalCell, ParameterUpdate, GetNextSingleOp, VirtualDatasetCellTriple, MicroBatchInterleaved, PipelineCell
+    WithEvalCell, ParameterUpdate, GetNextSingleOp, VirtualDatasetCellTriple, MicroBatchInterleaved, PipelineCell, \
+    GradAccumulationCell
 from mindspore.nn.wrap.loss_scale import TrainOneStepWithLossScaleCell,\
     DynamicLossScaleUpdateCell, FixedLossScaleUpdateCell
 from mindspore.nn.wrap.grad_reducer import DistributedGradReducer
@@ -42,5 +43,6 @@ __all__ = [
     "ParameterUpdate",
     "DynamicLossScaleUpdateCell",
     "FixedLossScaleUpdateCell",
-    "VirtualDatasetCellTriple"
+    "VirtualDatasetCellTriple",
+    "GradAccumulationCell"
     ]

mindspore/nn/wrap/cell_wrapper.py

@@ -673,7 +673,7 @@ class PipelineCell(Cell):
 
 class GradAccumulationCell(Cell):
     """
-    Wrap the network with Micro Batch.
+    Wrap the network with Micro Batch to enable the grad accumulation in semi_auto_parallel/auto_parallel mode.
 
     Args:
         network (Cell): The target network to wrap.
@@ -683,8 +683,11 @@ class GradAccumulationCell(Cell):
         ``Ascend`` ``GPU``
 
     Examples:
-        >>> net = Net()
-        >>> net = GradAccumulationCell(net, 4)
+        >>> import mindspore.nn as nn
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> net = nn.GradAccumulationCell(net, 4)
     """
     def __init__(self, network, micro_size):
         super(GradAccumulationCell, self).__init__(auto_prefix=False)

mindspore/numpy/math_ops.py

@@ -4285,7 +4285,7 @@ def argmin(a, axis=None):
 
     Examples:
         >>> import mindspore.numpy as np
-        >>> a = np.arange(10, 16).reshape(2, 3)
+        >>> a = np.arange(10, 16).reshape(2, 3).astype(np.float32)
         >>> print(np.argmin(a))
         0
         >>> print(np.argmin(a, axis=0))

mindspore/ops/__init__.py

@@ -34,6 +34,7 @@ from mindspore.ops.composite import *
 from mindspore.ops.operations import *
 from mindspore.ops.function import *
 from mindspore.ops.functional import *
+from mindspore.ops.silent_check import _silent_check
 
 __primitive__ = [
     "prim_attr_register", "Primitive", "PrimitiveWithInfer", "PrimitiveWithCheck", "signature"
@@ -48,3 +49,5 @@ __all__.extend(composite.__all__)
 __all__.extend(operations.__all__)
 __all__.extend(functional.__all__)
 __all__.extend(function.__all__)
+
+_silent_check()

mindspore/ops/_grad_experimental/grad_array_ops.py

@@ -36,8 +36,6 @@ from mindspore.ops.operations.array_ops import ScatterAddWithAxis
 from mindspore.ops.operations.array_ops import Expand
 from mindspore.ops.operations.array_ops import SegmentMean
 from mindspore.ops.operations.array_ops import AffineGrid
-from mindspore.ops.operations.array_ops import Im2Col
-from mindspore.ops.operations.array_ops import Col2Im
 from mindspore.ops.operations.array_ops import MaskedScatter
 from mindspore.ops.operations.array_ops import MaskedSelect
 from mindspore.ops.operations.array_ops import CountNonZero
@@ -360,35 +358,6 @@ def get_bprop_resize_nearest_neighbor_v2(self):
     return bprop
 
 
-@bprop_getters.register(Im2Col)
-def get_bprop_im2col(self):
-    """
-    Generate bprop for Im2Col
-
-    Im2Col, corresponding to torch's UnFold operator.
-    The Unfold operator has no `padding_mode` attribute,
-    and it's implementation corresponds to the mindspore
-    implementation with `padding_mode=CALCULATED` .
-    So, currently the bprop function of Im2Col only supports
-    the CALCULATED mode.
-    """
-    kernel_size = self.ksizes
-    dilation = self.dilations
-    stride = self.strides
-    padding = (self.pads[0], self.pads[-1])
-    col2im = Col2Im(kernel_size=kernel_size,
-                    dilation=dilation,
-                    stride=stride,
-                    padding=padding)
-
-    def bprop(x, out, dout):
-        x_shape = P.TensorShape()(x)[2:]
-        dx = col2im(dout, x_shape)
-        return (dx,)
-
-    return bprop
-
-
 @bprop_getters.register(P.ExtractVolumePatches)
 def get_bprop_extract_volume_patches(self):
     """Generate bprop for ExtractVolumePatches"""

mindspore/ops/_grad_experimental/grad_comm_ops.py

@@ -92,7 +92,8 @@ def get_bprop_send(self):
     """Generate bprop for Send."""
     shape = self.get_attr_dict()["shape"]
     dtype = self.get_attr_dict()["dtype"]
-    send_grad = Receive(self.sr_tag, self.rank, shape, dtype, self.group_back)
+    tag = self.get_attr_dict()["sr_tag"]
+    send_grad = Receive(tag, self.rank, shape, dtype, self.group_back)
     virtual_input = Tensor(0.0, dtype)
 
     def bprop(x, out, dout):
@@ -105,7 +106,8 @@ def get_bprop_send(self):
 @bprop_getters.register(Receive)
 def get_bprop_receive(self):
     """Generate bprop for Receive."""
-    receive_grad = Send(self.tag, self.rank, self.group_back)
+    tag = self.get_attr_dict()["sr_tag"]
+    receive_grad = Send(tag, self.rank, self.group_back)
     depend = P.Depend()
     cast = P.Cast()
     out_tensor = Tensor(0.0, mstype.float16)

mindspore/ops/_grad_experimental/grad_inner_ops.py

@@ -36,6 +36,14 @@ def get_bprop_parallel_resize_bilinear(self):
     return bprop
 
 
+@bprop_getters.register(P.inner_ops.GenerateEodMask)
+def get_bprop_generate_eod_mask(self):
+
+    def bprop(x, out, dout):
+        return dout, dout
+    return bprop
+
+
 @bprop_getters.register(inner.PsROIPooling)
 def get_bprop_ps_roi_pooling(self):
     """Grad definition for `PsROIPooling` operation."""

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -18,11 +18,13 @@
 import numpy as np
 import mindspore.numpy as mnp
 from mindspore.common import dtype as mstype
+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
+from mindspore.ops.operations.math_ops import Polar, 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
@@ -763,6 +765,7 @@ def get_bprop_fft_with_size(self):
     to_tensor_op = P.ScalarToTensor()
     type_op = P.DType()
     concat_op = P.Concat()
+    concat_op_last = P.Concat(axis=-1)
     ones_op = P.Ones()
     zeros_op = P.Zeros()
     real_op = P.Real()
@@ -794,8 +797,7 @@ def get_bprop_fft_with_size(self):
                                  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:])
+            irfft3d_ = FFTWithSize(signal_ndim=3, inverse=True, real=False, norm="backward", onesided=onesided)
             if inverse is False:
                 if onesided is True:
                     terms = 0
@@ -811,6 +813,7 @@ def get_bprop_fft_with_size(self):
                         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)
+                        dx = to_tensor_op(0.5, input_type) * (dx * rfft_offset_size + terms) * rfft_norm_offset
                     elif signal_ndim == 2:
                         dx = irfft2d_(dout)
                         arange_inner = mnp.arange(0, input_shape[-2], 1, input_type)
@@ -852,26 +855,27 @@ def get_bprop_fft_with_size(self):
                                                         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)
+                        dx = to_tensor_op(0.5, input_type) * (dx * rfft_offset_size + terms) * rfft_norm_offset
                     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
+                        zeros_shape = offset_shape[:-1] + (offset_shape[-1] - dout_shape[-1],)
+                        zeros_values = zeros_op(zeros_shape, input_type)
+                        zeros_padding = complex_op(zeros_values, zeros_values)
+                        dout = concat_op_last((dout, zeros_padding))
+                        dx = real_op(irfft3d_(dout)) * real_op(offset_size)
                 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))
+                    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,)
@@ -1017,3 +1021,19 @@ def get_bprop_tensor_add(self):
         return binop_grad_common(x, y, dout, dout)
 
     return bprop
+
+
+@bprop_getters.register(_MirrorSilentCheck)
+def get_bprop_mirror_silent_check(self):
+    """Grad definition for '_MirrorSilentCheck' op"""
+    silent_check = SilentCheck(self.min_steps, self.thresh_l1, self.coeff_l1, self.thresh_l2, self.coeff_l2)
+    out_tensor = Tensor([0.0], mstype.float32)
+
+    def bporp(x, pre_val, min_val, max_val, n_step, loss_scale, out, dout):
+        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)
+        return (dx, out_tensor, out_tensor, out_tensor, out_tensor, out_tensor)
+
+    return bporp

mindspore/ops/_op_impl/aicpu/__init__.py

@@ -60,6 +60,7 @@ from .init_data_set_queue import _init_data_set_queue_aicpu
 from .embedding_lookup import _embedding_lookup_aicpu
 from .padding import _padding_aicpu
 from .gather import _gather_aicpu
+from .generate_eod_mask import _generate_eod_mask_aicpu
 from .gather_grad import _gather_grad_aicpu
 from .gather_d_grad_v2 import _gather_d_grad_v2_aicpu
 from .gather_d import _gather_d_aicpu

mindspore/ops/_op_impl/aicpu/generate_eod_mask.py

@@ -0,0 +1,38 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""GenerateEodMask op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+generate_eod_mask_op_info = AiCPURegOp("GenerateEodMask") \
+    .fusion_type("OPAQUE") \
+    .attr("eod_token_id", "int") \
+    .attr("n_pos", "int") \
+    .attr("n_step", "listint") \
+    .attr("n_error_mode", "str") \
+    .input(0, "inputs_ids", "required") \
+    .output(0, "position_ids", "required") \
+    .dtype_format(DataType.U16_Default, DataType.U16_Default) \
+    .dtype_format(DataType.U32_Default, DataType.U32_Default) \
+    .dtype_format(DataType.U64_Default, DataType.U64_Default) \
+    .dtype_format(DataType.I32_Default, DataType.I32_Default) \
+    .dtype_format(DataType.I64_Default, DataType.I64_Default) \
+    .get_op_info()
+
+
+@op_info_register(generate_eod_mask_op_info)
+def _generate_eod_mask_aicpu():
+    """GenerateEodMask AiCPU register"""
+    return

mindspore/ops/_op_impl/aicpu/linear_sum_assignment.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Huawei Technologies Co., Ltd
+# Copyright 2023 Huawei Technologies Co., Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -24,11 +24,30 @@ lsap_op_info = AiCPURegOp("LinearSumAssignment") \
     .input(2, 'maximize', "required") \
     .output(0, "row_ind", "required") \
     .output(1, "col_ind", "required") \
-    .attr("cust_aicpu", "str") \
     .dtype_format(DataType.F64_Default, DataType.I64_Default,
                   DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
     .dtype_format(DataType.F32_Default, DataType.I64_Default,
                   DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.BOOL_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.I16_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.I32_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.I64_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.I8_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.U16_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.U32_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.U64_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
+    .dtype_format(DataType.U8_Default, DataType.I64_Default,
+                  DataType.BOOL_Default, DataType.I64_Default, DataType.I64_Default,) \
     .get_op_info()