mindspore 2.1.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.10__cp38-cp38-manylinux1_x86_64.whl

mindspore/ops/_op_impl/_custom_op/flash_attention/flash_attention_impl.py

@@ -15,7 +15,7 @@
 """The impl of flash attention"""
 from __future__ import absolute_import
 import mindspore.ops as ops
-from mindspore import dtype as mstype
+import mindspore.common.dtype as mstype
 from mindspore.ops import Custom
 from mindspore.ops import DataType
 from mindspore.ops import TBERegOp
@@ -39,31 +39,28 @@ cus_flash_atten_op_info = TBERegOp("FlashAttentionPrimitive") \
     .input(0, "query", False, "required", "all") \
     .input(1, "key", False, "required", "all") \
     .input(2, "value", False, "required", "all") \
-    .input(3, "dim_mask", False, "required", "all") \
-    .input(4, "attn_mask", False, "optional", "all") \
-    .input(5, "dropout_mask", False, "optional", "all") \
-    .input(6, "alibi_mask", False, "optional", "all") \
+    .input(3, "attn_mask", False, "optional", "all") \
+    .input(4, "dropout_mask", False, "optional", "all") \
+    .input(5, "alibi_mask", False, "optional", "all") \
     .output(0, "output", False, "required", "all") \
     .output(1, "rowsum", False, "required", "all") \
     .output(2, "rowmax", False, "required", "all") \
-    .dtype_format(DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.I8_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
                   DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
                   DataType.F16_Default,
                   DataType.F16_Default) \
-    .dtype_format(DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.I8_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
                   DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
                   DataType.F32_Default,
                   DataType.F16_Default) \
     .get_op_info()
@@ -88,41 +85,38 @@ cus_flash_atten_grad_op_info = TBERegOp("FlashAttentionGradPrimitive") \
     .input(4, "do", False, "required", "all") \
     .input(5, "rowsum", False, "required", "all") \
     .input(6, "rowmax", False, "required", "all") \
-    .input(7, "dim_mask", False, "required", "all") \
-    .input(8, "attn_mask", False, "optional", "all") \
-    .input(9, "dropout_mask", False, "optional", "all") \
-    .input(10, "alibi_mask", False, "optional", "all") \
+    .input(7, "attn_mask", False, "optional", "all") \
+    .input(8, "dropout_mask", False, "optional", "all") \
+    .input(9, "alibi_mask", False, "optional", "all") \
     .output(0, "dq", False, "required", "all") \
     .output(1, "dk", False, "required", "all") \
     .output(2, "dv", False, "required", "all") \
-    .dtype_format(DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.I8_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F32_Default,
-                  DataType.F32_Default,
-                  DataType.F32_Default) \
-    .dtype_format(DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F16_Default,
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ) \
+    .dtype_format(DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
+                  DataType.F16_FracNZ,
                   DataType.F32_Default,
                   DataType.F16_Default,
-                  DataType.I8_Default,
-                  DataType.F16_Default,
+                  DataType.F16_FracNZ,
                   DataType.F16_Default,
-                  DataType.F16_Default,
-                  DataType.F32_Default,
-                  DataType.F32_Default,
-                  DataType.F32_Default) \
+                  DataType.F16_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ,
+                  DataType.F32_FracNZ) \
     .get_op_info()
 
 
@@ -131,11 +125,11 @@ def get_flash_attention_grad(prev_block_num=65536, next_block_num=65536,
     """get flash attention grad"""
 
     def infer_shape(q_shape, k_shape, v_shape, o_shape, do_shape, l_shape, m_shape,
-                    dim_mask_shape, att_mask_shape, dropout_mask_shape, alibi_mask_shape):
+                    att_mask_shape, dropout_mask_shape, alibi_mask_shape):
         return q_shape, k_shape, v_shape
 
     def infer_dtype(q_dtype, k_dtype, v_dtype, o_dytpe, do_dtype, l_dtype, m_dtype,
-                    dim_mask_dtype, attn_mask_dtype, dropout_mask_dtype, alibi_mask_type):
+                    attn_mask_dtype, dropout_mask_dtype, alibi_mask_type):
         return mstype.float32, mstype.float32, mstype.float32
 
     fa_grad = Custom(flash_attention_grad, out_shape=infer_shape,
@@ -145,20 +139,20 @@ def get_flash_attention_grad(prev_block_num=65536, next_block_num=65536,
     fa_grad.add_prim_attr("high_precision", high_precision)
     fa_grad.add_prim_attr("tiling_stgy_name", tiling_stgy_name)
     fa_grad.init_prim_io_names(
-        inputs=["query", "key", "value", "output", "do", "rowsum", "rowmax", "dim_mask", "attn_mask", "dropout_mask",
+        inputs=["query", "key", "value", "output", "do", "rowsum", "rowmax", "attn_mask", "dropout_mask",
                 "alibi_mask"],
         outputs=["dq", "dk", "dv"]
     )
 
-    def bprop(query, key, value, dim_mask, attn_mask, dropout_mask, alibi_mask, out, douts):
+    def bprop(query, key, value, attn_mask, dropout_mask, alibi_mask, out, douts):
         output, rowsum, rowmax = out
         dout, _, _ = douts
-        dq, dk, dv = fa_grad(query, key, value, output, dout, rowsum, rowmax, dim_mask, attn_mask, dropout_mask,
+        dq, dk, dv = fa_grad(query, key, value, output, dout, rowsum, rowmax, attn_mask, dropout_mask,
                              alibi_mask)
         dq = ops.cast(dq, mstype.float16)
         dk = ops.cast(dk, mstype.float16)
         dv = ops.cast(dv, mstype.float16)
-        return dq, dk, dv, zeros_like(dim_mask), zeros_like(attn_mask), \
+        return dq, dk, dv, zeros_like(attn_mask), \
             zeros_like(dropout_mask), zeros_like(alibi_mask)
 
     return bprop
@@ -167,7 +161,7 @@ def get_flash_attention_grad(prev_block_num=65536, next_block_num=65536,
 def get_flash_attention(prev_block_num=65536, next_block_num=65536, tiling_stgy_name='sparse', high_precision=False):
     """get_flash_attention"""
 
-    def infer_shape(q_shape, k_shape, v_shape, dim_mask_shape, attn_mask_shape=None,
+    def infer_shape(q_shape, k_shape, v_shape, attn_mask_shape=None,
                     dropout_mask_shape=None, alibi_mask_shape=None):
         """infer shape"""
         batch, hidden_size, seq_len, _ = q_shape
@@ -175,7 +169,7 @@ def get_flash_attention(prev_block_num=65536, next_block_num=65536, tiling_stgy_
         m_shape = (batch, hidden_size, seq_len)
         return q_shape, l_shape, m_shape
 
-    def infer_dtype(q_dtype, k_dtype, v_dtype, dim_mask_dtype, attn_mask_dtype=None,
+    def infer_dtype(q_dtype, k_dtype, v_dtype, attn_mask_dtype=None,
                     dropout_mask_dtype=None, alibi_mask_type=None):
         """infer type"""
         l_dtype = mstype.float16
@@ -192,7 +186,7 @@ def get_flash_attention(prev_block_num=65536, next_block_num=65536, tiling_stgy_
     fa_forward.add_prim_attr("high_precision", high_precision)
     fa_forward.add_prim_attr("tiling_stgy_name", tiling_stgy_name)
     fa_forward.init_prim_io_names(
-        inputs=["query", "key", "value", "dim_mask", "attn_mask", "dropout_mask", "alibi_mask"],
+        inputs=["query", "key", "value", "attn_mask", "dropout_mask", "alibi_mask"],
         outputs=["output", "rowsum", "rowmax"]
     )
 

mindspore/ops/_op_impl/_custom_op/flash_attention/tik_ops_utils.py

@@ -19,7 +19,6 @@ from mindspore.ops._op_impl._custom_op.flash_attention.constants import DTYPE_SI
 from mindspore.ops._op_impl._custom_op.flash_attention.constants import FP16
 from mindspore.ops._op_impl._custom_op.flash_attention.constants import FP32
 from mindspore.ops._op_impl._custom_op.flash_attention.constants import L0C
-from mindspore.ops._op_impl._custom_op.flash_attention.constants import L1
 from mindspore.ops._op_impl._custom_op.flash_attention.constants import UB
 
 
@@ -179,7 +178,7 @@ class TikOpsUtils:
     def broadcast(self, vec_ub, shape):
         """ broadcast a vector to a matrix
         :param vec_ub: a tensor in UB with shape of (M,), and dtype is float16
-        :param shape: the target shape, a tuple with value (M, N)，M and N are integer multiples of 16
+        :param shape: the target shape, a tuple with value (M, N), M and N are integer multiples of 16
         :return: a tensor in UB with shape of (M, N)
         """
         M, N = shape
@@ -321,27 +320,16 @@ class TikOpsUtils:
                 )
         return vec_rec_ub
 
-    def row_sum_cube_impl(self, matrix_l1_K1MK0_ed, rowsum_ub, m, k, precision_type):
+    def row_sum_cube_impl(self, matrix_l1_K1MK0_ed, right_all_one_matrix_l1, rowsum_ub, m, k, precision_type):
         """用cube实现矩阵行和：右乘一个shape=(n,1)全一矩阵
         :param matrix_l1_K1MK0_ed: input tensor with shape (K1, M, K0)
-        :param rowsum_ub: output tensor stores the row sum of input tensor.
+        :param right_all_one_matrix_l1: input tensor with shape (K, 16)
+        :param rowsum_ub: output tensor stores the row sum of input tensor
         :param m: actual tensor height
         :param k: actual tensor width
         :return: row sum of the output tensor
         """
         K1, M, K0 = matrix_l1_K1MK0_ed.shape
-        K = K1 * K0
-
-        # 构造全一右矩阵，由于cube无法处理shape=(n, 1)，所以shape=(n, 16)，全一矩阵不需分形
-        right_all_one_matrix_ub = self.tik_instance.Tensor(
-            FP16, (K, 16), name="right_all_one_matrix_ub", scope=UB
-        )
-        self.tik_instance.h_duplicate(right_all_one_matrix_ub, 1.0)
-        right_all_one_matrix_l1 = self.tik_instance.Tensor(
-            FP16, (K1 * K0, 16), name="right_all_one_matrix_l1", scope=L1
-        )
-        self.cont_data_mv_1_bust(dst=right_all_one_matrix_l1, src=right_all_one_matrix_ub, burst=K)
-
         # 调用matmul实现rowsum，结果shape=(m, 16)，取每行的第一个数
         with self.tik_instance.new_stmt_scope(disable_sync=False):
             row_sum_ub_N1MN0 = self.matmul_compute(matrix_l1_K1MK0_ed, right_all_one_matrix_l1, m, k, 16,
@@ -352,6 +340,7 @@ class TikOpsUtils:
                     cur_row_sum = self.tik_instance.Scalar(FP32, init_value=row_sum_ub_MN_ed[idx, 0])
                     rowsum_ub[idx].set_as(cur_row_sum)
             else:
+                # row_sum_ub_MN_ed 先转置，然后取一行, 替换原来按行操作: lij_ub[i].set_as(row_sum_ub_MN_ed[i, 0])
                 row_sum_ub_trans = self.tik_instance.Tensor(FP16, (16, M), name="row_sum_ub_trans", scope=UB)
                 row_sum_ub_trans = self.transpose_matrix(row_sum_ub_MN_ed, row_sum_ub_trans, M, True)
                 self.cont_data_mv_1_bust(dst=rowsum_ub, src=row_sum_ub_trans, burst=M // 16)
@@ -409,7 +398,7 @@ class TikOpsUtils:
             offset = vec_len - a_burst_num
             last_blk_ub = self.tik_instance.Tensor(FP16, (a_burst_num,), name="last_blk_ub", scope=UB)
             self.cont_data_mv_1_bust(dst=last_blk_ub, src=src_tensor[gm_offset + offset], burst=1)
-            with self.tik_instance.for_range(0, a_burst_num) as idx:  # offset非32bytes对齐，无法用datamove
+            with self.tik_instance.for_range(0, a_burst_num) as idx:  # offset非32bytes对齐, 无法用datamove
                 dst_tensor[offset + idx].set_as(last_blk_ub[idx])
 
     def move_vector_from_ub_to_gm(self, dst_tensor, src_tensor, gm_offset, block_h):

mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/wukong_tiling.py

@@ -29,7 +29,7 @@ class WukongTiling(TilingStrategy):
         反向的空间分布待详细分析
         N = (4096, 1024, 256, 64) 或 77
         Nq = (4096, 1024, 256, 64)
-        d = dv = (40, 80, 160， 160)
+        d = dv = (40, 80, 160,  160)
         """
         if self.N <= 77:  # [77, 64]
             # cross-attention or self-attention of (64, 64, 160)

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/aicpu/eps.py

@@ -0,0 +1,32 @@
+# 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.
+# 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.
+# ============================================================================
+
+"""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/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") \

mindspore/ops/_op_impl/aicpu/parameterized_truncated_normal.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.
@@ -22,21 +22,29 @@ parameterized_truncated_normal_op_info = AiCPURegOp("ParameterizedTruncatedNorma
     .input(2, "stdevs", "required") \
     .input(3, "min", "required") \
     .input(4, "max", "required") \
+    .input(5, "counts", "required") \
+    .input(6, "states", "required") \
     .output(0, "y", "required") \
     .attr("seed", "int")\
     .attr("seed2", "int")\
     .dtype_format(DataType.I32_Default, DataType.F16_Default, DataType.F16_Default,
-                  DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \
+                  DataType.F16_Default, DataType.F16_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.F16_Default) \
     .dtype_format(DataType.I32_Default, DataType.F32_Default, DataType.F32_Default,
-                  DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \
+                  DataType.F32_Default, DataType.F32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.F32_Default) \
     .dtype_format(DataType.I32_Default, DataType.F64_Default, DataType.F64_Default,
-                  DataType.F64_Default, DataType.F64_Default, DataType.F64_Default) \
+                  DataType.F64_Default, DataType.F64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.F64_Default) \
     .dtype_format(DataType.I64_Default, DataType.F16_Default, DataType.F16_Default,
-                  DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \
+                  DataType.F16_Default, DataType.F16_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.F16_Default) \
     .dtype_format(DataType.I64_Default, DataType.F32_Default, DataType.F32_Default,
-                  DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \
+                  DataType.F32_Default, DataType.F32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.F32_Default) \
     .dtype_format(DataType.I64_Default, DataType.F64_Default, DataType.F64_Default,
-                  DataType.F64_Default, DataType.F64_Default, DataType.F64_Default) \
+                  DataType.F64_Default, DataType.F64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.F64_Default) \
     .get_op_info()
 
 

mindspore/ops/_op_impl/aicpu/random_categorical.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.
@@ -21,25 +21,45 @@ random_categorical_op_info = AiCPURegOp("RandomCategorical") \
     .input(0, "logits", "required") \
     .input(1, "num_sample", "required") \
     .input(2, "seed", "required") \
+    .input(3, "counts", "required") \
+    .input(4, "states", "required") \
     .output(0, "output", "required") \
-    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default, DataType.I16_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I16_Default) \
-    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default, DataType.I16_Default) \
-    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default, DataType.I64_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I64_Default) \
-    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default, DataType.I64_Default) \
-    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default, DataType.I16_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default, DataType.I16_Default) \
-    .dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default, DataType.I16_Default) \
-    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default, DataType.I32_Default) \
-    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
-    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
-    .dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.I32_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I64_Default, DataType.I64_Default, DataType.U64_Default,
+                  DataType.U64_Default, DataType.I64_Default) \
     .get_op_info()
 
 @op_info_register(random_categorical_op_info)