mindspore 2.1.0__cp38-none-any.whl → 2.2.11__cp38-none-any.whl

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -34,14 +34,12 @@ 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 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 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.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 _primexpr
@@ -137,7 +135,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 == mstype.complex64 or matrix.dtype == mstype.complex128:
+        if matrix.dtype in (mstype.complex64, mstype.complex128):
             grad_rhs_temp = _adjoint(grad_rhs)
             out_temp = _adjoint(out)
         else:
@@ -158,14 +156,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 == mstype.complex64 or grad_matrix.dtype == mstype.complex128:
+            if grad_matrix.dtype in (mstype.complex64, 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:
+            if grad_matrix.dtype in (mstype.complex64, 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 +219,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
 
 
@@ -424,115 +422,6 @@ def get_bprop_polar(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(TridiagonalSolve)
 def get_bprop_tridiagonalsolve(self):
     """Grad definition for 'TridiagonalSolve' operation"""
@@ -758,20 +647,21 @@ def _fft_rank_offset(norm_shape, rank):
 @_primexpr
 def _fft_with_size_back_norm(norm_shape, norm, inverse, rank):
     """generate reverse term for fft_with_size"""
+    norm_ = None
     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:
+            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:
         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
+            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
     return norm_
 
 
@@ -781,9 +671,9 @@ def _rfft_norm(norm_shape, norm, rank):
     norm_ = 1.0
     if norm == "forward":
         norm_ = 1 / _fft_rank_offset(norm_shape, rank)
-    if norm == "backward":
-        norm_ = 1
-    if norm == "ortho":
+    elif norm == "backward":
+        norm_ = 1.0
+    elif norm == "ortho":
         norm_ = 1 / np.sqrt(_fft_rank_offset(norm_shape, rank))
     return norm_
 
@@ -1127,73 +1017,3 @@ def get_bprop_tensor_add(self):
         return binop_grad_common(x, y, dout, dout)
 
     return bprop
-
-
-@bprop_getters.register(P.BitwiseAnd)
-def get_bprop_bitwiseand(self):
-    """Grad definition for `BitwiseAnd` operation."""
-
-    def bprop(x, y, out, dout):
-        return zeros_like(x), zeros_like(y)
-
-    return bprop
-
-
-@bprop_getters.register(P.BitwiseOr)
-def get_bprop_bitwiseor(self):
-    """Grad definition for `BitwiseOr` operation."""
-
-    def bprop(x, y, out, dout):
-        return zeros_like(x), zeros_like(y)
-
-    return bprop
-
-
-@bprop_getters.register(P.BitwiseXor)
-def get_bprop_bitwisexor(self):
-    """Grad definition for `BitwiseXor` operation."""
-
-    def bprop(x, y, out, dout):
-        return zeros_like(x), zeros_like(y)
-
-    return bprop
-
-
-@bprop_getters.register(P.InplaceUpdate)
-def get_bprop_inplace_update(self):
-    """Grad definition for `InplaceUpdate` operation."""
-
-    def bprop(x, v, out, dout):
-        return zeros_like(x), zeros_like(v)
-
-    return bprop
-
-
-@bprop_getters.register(P.InplaceUpdateV2)
-def get_bprop_inplace_update_v2(self):
-    """Grad definition for `InplaceUpdateV2` operation."""
-
-    def bprop(x, indices, v, out, dout):
-        return zeros_like(x), zeros_like(indices), zeros_like(v)
-
-    return bprop
-
-
-@bprop_getters.register(P.InplaceSub)
-def get_bprop_inplace_sub(self):
-    """Grad definition for `InplaceSub` operation."""
-
-    def bprop(x, input_v, out, dout):
-        return zeros_like(x), zeros_like(input_v)
-
-    return bprop
-
-
-@bprop_getters.register(P.InplaceAdd)
-def get_bprop_inplace_add(self):
-    """Grad definition for `InplaceAdd` operation."""
-
-    def bprop(x, input_v, out, dout):
-        return zeros_like(x), zeros_like(input_v)
-
-    return bprop

mindspore/ops/_grad_experimental/grad_sparse.py

@@ -1,3 +1,18 @@
+# Copyright 2022 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.
+# ============================================================================
+
 """Define the grad rules of math related operations."""
 
 from mindspore.ops import functional as F

mindspore/ops/_grad_experimental/grad_sparse_ops.py

@@ -358,10 +358,10 @@ def get_bprop_ragged_tensor_to_sparse(self):
                 split.append(zeros_like(i))
             all_d = (split, ragged_values_grad)
             return all_d
-        split = ()
+        split_ = ()
         for i in enumerate(rt_nested_splits):
-            split = split + (zeros_like(i),)
-        all_d = (split, ragged_values_grad)
+            split_ = split_ + (zeros_like(i),)
+        all_d = (split_, ragged_values_grad)
         return all_d
 
     return bprop

mindspore/ops/_op_impl/_custom_op/dsd_back_impl.py

@@ -299,7 +299,7 @@ def dsdbpropimpl(w1_gm, w2_gm, v_gm, a_gm, d_a_gm, d_w1_gm={}, d_w2_gm={}, d_v_g
                                                            global_size // 16, 16, 16),
                                                name='v_global_l0b', scope=tik.scope_cb)
 
-                # d_w_global，小z大n
+                # d_w_global, 小z大n
                 d_w_global_l0c = tik_inst.Tensor('float32', (global_size // 16, head_size // (16 * ub_time), 16, 16),
                                                  name='d_w_global_l0c', scope=tik.scope_cc)
                 d_w_global_ub = tik_inst.Tensor('float16', (global_size // 16,

mindspore/ops/_op_impl/aicpu/__init__.py

@@ -108,6 +108,7 @@ from .search_sorted import _search_sorted_aicpu
 from .stack import _stack_aicpu
 from .unstack import _unstack_aicpu
 from .unsorted_segment_sum import _unsorted_segment_sum_aicpu
+from .unsorted_segment_prod import _unsorted_segment_prod_aicpu
 from .addcmul import _addcmul_aicpu
 from .uniform_candidate_sampler import _uniform_candidate_sampler_aicpu
 from .log_uniform_candidate_sampler import _log_uniform_candidate_sampler_aicpu
@@ -145,6 +146,7 @@ from .upsample_trilinear_3d import _upsample_trilinear_3d_aicpu
 from .upsample_trilinear_3d_grad import _upsample_trilinear_3d_grad_aicpu
 from .upper_bound import _upper_bound_aicpu
 from .cache_swap_table import _cache_swap_table_aicpu
+from .uniform import _uniform_aicpu
 from .uniform_int import _uniform_int_aicpu
 from .uniform_real import _uniform_real_aicpu
 from .standard_laplace import _standard_laplace_aicpu
@@ -156,12 +158,13 @@ from .fused_sparse_adam import _fused_sparse_adam_aicpu
 from .fused_sparse_lazy_adam import _fused_sparse_lazy_adam_aicpu
 from .fused_sparse_ftrl import _fused_sparse_ftrl_aicpu
 from .sparse_fill_empty_rows_grad import _sparse_fill_empty_rows_grad_aicpu
+from .sparse_reorder import _sparse_reorder_aicpu
 from .sparse_reshape import _sparse_reshape_aicpu
 from .sparse_segment_sqrt_n_grad import _sparse_segment_sqrt_n_grad_aicpu
 from .sparse_segment_sum import _sparse_segment_sum_aicpu
 from .sparse_segment_sum_with_num_segments import _sparse_segment_sum_with_num_segments_aicpu
 from .sparse_softmax_cross_entropy_with_logits_v2 import _sparse_softmax_cross_entropy_with_logits_v2_aicpu
-from .sparsesparsemaximum import _sparsesparsemaximum_aicpu
+from .sparse_sparse_maximum import _sparse_sparse_maximum_aicpu
 from .split import _split_aicpu
 from .transpose import _transpose_aicpu
 from .tril_indices import _tril_indices_aicpu
@@ -205,6 +208,7 @@ from .environ_get import _environ_get_aicpu
 from .environ_destroy_all import _environ_destroy_all_aicpu
 from .cross import _cross_aicpu
 from .check_numerics import _check_numerics_aicpu
+from .cummax import _cummax_aicpu
 from .cumsum import _cumsum_aicpu
 from .round import _round_aicpu
 from .stft import _stft_aicpu
@@ -229,6 +233,7 @@ from .scatter_nd_update import _scatter_nd_update_aicpu
 from .scatter_nd_max import _scatter_nd_max_aicpu
 from .conj import _conj_aicpu
 from .scatter_nd_min import _scatter_nd_min_aicpu
+from .scatter_add_with_axis import _scatter_add_with_axis_aicpu
 from .compare_and_bitpack import _compare_and_bitpack_aicpu
 from .addcdiv import _addcdiv_aicpu
 from .unique_consecutive import _unique_consecutive_aicpu
@@ -241,8 +246,8 @@ from .reservoir_replay_buffer import _rrb_push_op_cpu
 from .reservoir_replay_buffer import _rrb_sample_op_cpu
 from .reservoir_replay_buffer import _rrb_destroy_op_cpu
 from .concat_offset import _concat_offset_aicpu
-from .concat_offset_v1 import _concat_offset_v1_aicpu
 from .range import _range_aicpu
+from .range_v2 import _range_v2_aicpu
 from .slice_grad import _slice_grad_aicpu
 from .median import _median_aicpu
 from .median_grad import _median_grad_aicpu
@@ -272,6 +277,7 @@ from .complex import _complex_aicpu
 from .complex_abs import _complex_abs_aicpu
 from .concat import _concat_aicpu
 from .cos import _cos_aicpu
+from .count_nonzero import _count_nonzero_aicpu
 from .csr_sparse_matrix_to_dense import _csr_sparse_matrix_to_dense_aicpu
 from .cumprod import _cumprod_aicpu
 from .exp import _exp_aicpu
@@ -340,6 +346,7 @@ from .hypot import _hypot_aicpu
 from .identity_n import _identity_n_aicpu
 from .index_fill import _index_fill_aicpu
 from .index_put import _index_put_aicpu
+from .inplace_index_add import _inplace_index_add_aicpu
 from .kldivloss import _kldiv_loss_aicpu
 from .kldivlossgrad import _kldiv_loss_grad_aicpu
 from .lcm import _lcm_aicpu
@@ -400,6 +407,9 @@ from .non_deterministic_ints import _non_deterministic_ints_aicpu
 from .pow import _pow_aicpu
 from .real import _real_aicpu
 from .resize_area import _resize_area_aicpu
+from .segment_mean import _segment_mean_aicpu
+from .segment_min import _segment_min_aicpu
+from .segment_prod import _segment_prod_aicpu
 from .segment_sum import _segment_sum_aicpu
 from .set_size import _set_size_aicpu
 from .slice import _slice_aicpu
@@ -411,6 +421,7 @@ from .sparse_tensor_dense_mat_mul import _sparse_tensor_dense_mat_mul_aicpu
 from .trace import _trace_aicpu
 from .tracegrad import _tracegrad_aicpu
 from .tridiagonal_solve import _tridiagonal_solve_aicpu
+from .tridiagonal_matmul import _tridiagonal_matmul_aicpu
 from .truncated_normal import _truncated_normal_aicpu
 from .glu import _glu_aicpu
 from .deformable_offsets import _deformable_offsets_aicpu
@@ -426,3 +437,4 @@ from .sequence_concat import _sequence_concat_aicpu
 from .sequence_stack import _sequence_stack_aicpu
 from .affine_grid import _affine_grid_aicpu
 from .depth_to_space import _depth_to_space_aicpu
+from .eps import _eps_aicpu

mindspore/ops/_op_impl/aicpu/add.py

@@ -29,9 +29,9 @@ add_op_info = AiCPURegOp("Add") \
     .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
     .dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
     .dtype_format(DataType.U8_Default, DataType.U8_Default, DataType.U8_Default) \
-    .dtype_format(DataType.U16_Default, DataType.I16_Default, DataType.I16_Default) \
-    .dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.I32_Default) \
-    .dtype_format(DataType.U64_Default, DataType.I64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U16_Default, DataType.U16_Default, DataType.U16_Default) \
+    .dtype_format(DataType.U32_Default, DataType.U32_Default, DataType.U32_Default) \
+    .dtype_format(DataType.U64_Default, DataType.U64_Default, DataType.U64_Default) \
     .dtype_format(DataType.C64_Default, DataType.C64_Default, DataType.C64_Default) \
     .dtype_format(DataType.C128_Default, DataType.C128_Default, DataType.C128_Default) \
     .get_op_info()

mindspore/ops/_op_impl/aicpu/bias_add_grad.py

@@ -31,7 +31,6 @@ bias_add_grad_op_info = AiCPURegOp("BiasAddGrad") \
     .dtype_format(DataType.I64_Default, DataType.I64_Default) \
     .dtype_format(DataType.F16_Default, DataType.F16_Default) \
     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
-    .dtype_format(DataType.F64_Default, DataType.F64_Default) \
     .dtype_format(DataType.C64_Default, DataType.C64_Default) \
     .dtype_format(DataType.C128_Default, DataType.C128_Default) \
     .get_op_info()

mindspore/ops/_op_impl/aicpu/count_nonzero.py

@@ -0,0 +1,43 @@
+# Copyright 2022 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.
+# ============================================================================
+
+"""CountNonZero op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+count_nonzero_op_info = AiCPURegOp("CountNonZero") \
+    .fusion_type("OPAQUE") \
+    .input(0, "x", "required") \
+    .output(0, "y", "required") \
+    .attr("dims", "listInt")\
+    .dtype_format(DataType.I8_Default, DataType.I64_Default) \
+    .dtype_format(DataType.I16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.I32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.I64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U8_Default, DataType.I64_Default)  \
+    .dtype_format(DataType.U16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I64_Default) \
+    .get_op_info()
+
+
+@op_info_register(count_nonzero_op_info)
+def _count_nonzero_aicpu():
+    """CountNonZero AiCPU register"""
+    return

mindspore/ops/_op_impl/{_custom_op/flash_attention/constants.py → aicpu/eps.py}

@@ -1,4 +1,4 @@
-# Copyright 2023 Huawei Technologies Co., Ltd
+# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,30 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""define constants"""
-from tbe import tik
 
-BLOCK_NUM = 16
-FP16 = "float16"
-INT8 = "int8"
-INT32 = "int32"
-FP32 = "float32"
-REPEAT_SZ = 128
-BLK_STRIDE = 1
-REPEAT_STRIDE = 8
-TRANS_CUBE_TGT = 8
-FP16_MIN_VAL = -65504.0
-MASK_FILL_VALUE = -10000.0
-GM = tik.scope_gm
-L1 = tik.scope_cbuf
-L1OUT = tik.scope_cbuf_out
-UB = tik.scope_ubuf
-L0A = tik.scope_ca
-L0B = tik.scope_cb
-L0C = tik.scope_cc
-DTYPE_SIZE = {
-    "int8": 1,
-    "float16": 2,
-    "int16": 2,
-    "float32": 4,
-}
+"""Eps op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+eps_op_info = AiCPURegOp("Eps") \
+    .fusion_type("OPAQUE") \
+    .input(0, "x", "required") \
+    .output(0, "y", "required") \
+    .dtype_format(DataType.F16_Default, DataType.F16_Default)   \
+    .dtype_format(DataType.F32_Default, DataType.F32_Default)   \
+    .dtype_format(DataType.F64_Default, DataType.F64_Default)   \
+    .get_op_info()
+
+
+@op_info_register(eps_op_info)
+def _eps_aicpu():
+    """Eps AiCPU register"""
+    return

mindspore/ops/_op_impl/aicpu/gamma.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 
-"""RandomGamma op"""
+"""Gamma op"""
 from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
 
 gamma_op_info = AiCPURegOp("Gamma") \
@@ -32,5 +32,5 @@ gamma_op_info = AiCPURegOp("Gamma") \
 
 @op_info_register(gamma_op_info)
 def _gamma_aicpu():
-    """RandomGamma AiCPU register"""
+    """Gamma 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()
 
 

mindspore/ops/_op_impl/aicpu/log_uniform_candidate_sampler.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Huawei Technologies Co., Ltd
+# Copyright 2020-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.
@@ -18,15 +18,18 @@ from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataTyp
 log_uniform_candidate_sampler_op_info = AiCPURegOp("LogUniformCandidateSampler") \
     .fusion_type("OPAQUE") \
     .input(0, "true_classes", "required") \
+    .input(1, "counts", "required") \
+    .input(2, "states", "required") \
     .output(0, "sampled_candidates", "required") \
     .output(1, "true_expected_count", "required") \
-    .output(2, "true_expected_count", "required") \
+    .output(2, "sampled_expected_count", "required") \
     .attr("num_true", "int") \
     .attr("num_sampled", "int") \
     .attr("unique", "bool") \
     .attr("range_max", "int") \
     .attr("seed", "int") \
-    .dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.F32_Default, DataType.F32_Default) \
+    .dtype_format(DataType.I64_Default, DataType.U64_Default, DataType.U64_Default, DataType.I64_Default,
+                  DataType.F32_Default, DataType.F32_Default) \
     .get_op_info()
 
 

mindspore/ops/_op_impl/aicpu/lu_unpack_grad.py

@@ -19,7 +19,6 @@ from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataTyp
 lu_unpack_grad_op_info = AiCPURegOp("LuUnpackGrad") \
     .fusion_type("OPAQUE") \
     .attr("L_grad_flag", "bool") \
-    .attr("L_grad_flag", "bool") \
     .input(0, "L_grad", "required") \
     .input(1, "U_grad", "required") \
     .input(2, "LU_data", "required") \

mindspore/ops/_op_impl/aicpu/multinomial.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Huawei Technologies Co., Ltd
+# Copyright 2022-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.
@@ -20,8 +20,8 @@ multinomial_op_info = AiCPURegOp("Multinomial") \
     .fusion_type("OPAQUE") \
     .input(0, "input", "required") \
     .input(1, "num_sample", "required") \
-    .input(2, "count", "required") \
-    .input(3, "state", "required") \
+    .input(2, "counts", "required") \
+    .input(3, "states", "required") \
     .output(0, "output", "required") \
     .attr("dtype", "Type") \
     .attr("seed", "int") \