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

mindspore/ops/__init__.py

@@ -29,13 +29,14 @@ from mindspore.ops.vm_impl_registry import get_vm_impl_fn, vm_impl_registry
 from mindspore.ops.op_info_register import op_info_register, custom_info_register, AkgGpuRegOp, AkgAscendRegOp, \
     AiCPURegOp, TBERegOp, CpuRegOp, CustomRegOp, DataType
 from mindspore.ops.primitive import constexpr
-from mindspore.ops import composite, operations, functional, function, auto_generate, extend
+from mindspore.ops import composite, operations, functional, function
 from mindspore.ops import signature
+from mindspore.ops.auto_generate import cpp_create_prim_instance_helper, gen_arg_dtype_cast, gen_arg_handler, \
+    gen_extend_func, gen_ops_def, gen_ops_prim, pyboost_inner_prim
 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", "prim_arg_register", "Primitive", "PrimitiveWithInfer", "PrimitiveWithCheck", "signature"
@@ -44,11 +45,12 @@ __primitive__ = [
 __all__ = ["get_vm_impl_fn", "vm_impl_registry",
            "op_info_register", "custom_info_register", "AkgGpuRegOp", "AkgAscendRegOp", "AiCPURegOp", "TBERegOp",
            "CpuRegOp", "CustomRegOp", "DataType",
-           "constexpr", "reshard"]
+           "constexpr", "reshard",
+           "cpp_create_prim_instance_helper", "gen_arg_dtype_cast", "gen_arg_handler", "gen_extend_func", "gen_ops_def",
+           "gen_ops_prim", "pyboost_inner_prim"]
 __all__.extend(__primitive__)
 __all__.extend(composite.__all__)
 __all__.extend(operations.__all__)
 __all__.extend(functional.__all__)
 __all__.extend(function.__all__)
 __all__.extend(auto_generate.__all__)
-_silent_check()

mindspore/ops/_grad_experimental/grad_array_ops.py

@@ -38,7 +38,6 @@ from mindspore.ops.operations.array_ops import SegmentMean
 from mindspore.ops.operations.array_ops import AffineGrid
 from mindspore.ops.operations.array_ops import MaskedScatter
 from mindspore.ops.operations.array_ops import MaskedSelect
-from mindspore.ops.operations.array_ops import CountNonZero
 from mindspore.ops.operations.random_ops import LogNormalReverse
 from mindspore.ops.operations.random_ops import ParameterizedTruncatedNormal
 from mindspore.ops.operations import _inner_ops as inner
@@ -125,16 +124,6 @@ def get_bprop_masked_scatter(self):
     return bprop
 
 
-@bprop_getters.register(CountNonZero)
-def get_bprop_countnonzero(self):
-    """Grad definition for CountNonZero"""
-
-    def bprop(x, out, dout):
-        return (zeros_like(x),)
-
-    return bprop
-
-
 @bprop_getters.register(Mvlgamma)
 def get_bprop_mvlgamma(self):
     """Grad definition for Mvlgamma"""

mindspore/ops/_grad_experimental/grad_comm_ops.py

@@ -31,9 +31,11 @@ from mindspore.ops.operations.comm_ops import (AllGather, _MiniStepAllGather, _H
                                                _GetTensorSlice, _MirrorOperator, _MirrorMiniStepOperator, ReduceOp,
                                                ReduceScatter, _HostReduceScatter, _VirtualDiv, _VirtualAdd, _AllSwap,
                                                _VirtualAssignAdd, _VirtualAccuGrad, _MirrorMicroStepOperator,
-                                               _MicroStepAllGather, Reduce, CollectiveGather, CollectiveScatter)
+                                               _MicroStepAllGather, Reduce, CollectiveGather, CollectiveScatter,
+                                               _VirtualAssignKvCache)
 from mindspore.ops._grad_experimental.grad_base import bprop_getters
 from mindspore.ops.operations import _grad_ops as G
+import mindspore as ms
 
 
 @bprop_getters.register(AllReduce)
@@ -95,6 +97,12 @@ def get_bprop_send(self):
     dtype = self.get_attr_dict()["dtype"]
     tag = self.get_attr_dict()["sr_tag"]
     send_grad = Receive(tag, self.rank, shape, dtype, self.group_back)
+    if "dst_global_rank" in self.get_attr_dict():
+        dst_global_rank = self.get_attr_dict().get("dst_global_rank")
+        send_grad.add_prim_attr("src_global_rank", dst_global_rank)
+    if "RING_ATTENTION_INDEX" in self.get_attr_dict():
+        ringattention = self.get_attr_dict().get("RING_ATTENTION_INDEX")
+        send_grad.add_prim_attr("RING_ATTENTION_INDEX", ringattention)
     virtual_input = Tensor(0.0, dtype)
 
     def bprop(x, out, dout):
@@ -108,8 +116,16 @@ def get_bprop_send(self):
 def get_bprop_receive(self):
     """Generate bprop for Receive."""
     tag = self.get_attr_dict()["sr_tag"]
+    flash_tag = self.get_attr_dict().get("flash_tag")
     receive_grad = Send(tag, self.rank, self.group_back)
-    receive_grad.add_prim_attr("shape", self.shape)
+    shape = self.get_attr_dict()["shape"]
+    receive_grad.add_prim_attr("shape", shape)
+    if "src_global_rank" in self.get_attr_dict():
+        src_global_rank = self.get_attr_dict().get("src_global_rank")
+        receive_grad.add_prim_attr("dst_global_rank", src_global_rank)
+    if "RING_ATTENTION_INDEX" in self.get_attr_dict():
+        ringattention = self.get_attr_dict().get("RING_ATTENTION_INDEX")
+        receive_grad.add_prim_attr("RING_ATTENTION_INDEX", ringattention)
     depend = P.Depend()
     cast = P.Cast()
     out_tensor = Tensor(0.0, mstype.float16)
@@ -117,7 +133,7 @@ def get_bprop_receive(self):
 
     def bprop(x, out, dout):
         send_out = receive_grad(dout)
-        if is_opt_shard:
+        if is_opt_shard or (flash_tag == "True"):
             dx = depend(F.zeros_like(x), send_out)
         else:
             dx = depend(cast(out_tensor, F.dtype(x)), send_out)
@@ -164,6 +180,24 @@ def get_bprop_virtual_assign_add(self):
     return bprop
 
 
+@bprop_getters.register(_VirtualAssignKvCache)
+def get_bprop_virtual_assign_kv_cache(self):
+    """Generate bprop for VirtualAssignAdd."""
+    assign = P.Assign()
+    cast = P.Cast()
+    dtype = P.DType()
+    out_tensor = Tensor(0.0, mstype.float16)
+
+    def bprop(x, y, seq_chunk, out, dout):
+        dout_update = dout + y
+        kv_equal = F.equal(seq_chunk, 0)
+        update_kv = F.select(kv_equal, F.broadcast_to(cast(out_tensor, dtype(y)), F.shape(y)), dout_update)
+        return F.depend((dout_update, cast(out_tensor, dtype(y)),
+                         cast(out_tensor, dtype(seq_chunk))), assign(y, update_kv))
+
+    return bprop
+
+
 @bprop_getters.register(_VirtualAccuGrad)
 def get_bprop_virtual_accu_grad(self):
     """Generate bprop for VirtualAccuGrad."""
@@ -186,6 +220,9 @@ def get_bprop_mirror_micro_step_operator(self):
     group = self.group
     dev_num = self.dev_num
     mean_flag = self.mean_flag
+    param_name = " "
+    if 'mirror_user_id' in self.get_attr_dict():
+        param_name = self.get_attr_dict()['mirror_user_id']
     scale = 1 / dev_num
 
     all_reduce = AllReduce(group=group)
@@ -196,7 +233,6 @@ def get_bprop_mirror_micro_step_operator(self):
     if hasattr(self, 'parameter'):
         parameter = self.parameter
         all_reduce.add_prim_attr("parameter", parameter)
-
     if self.instance_name:
         instance_name = "grad_mirror" + self.instance_name
         all_reduce.set_prim_instance_name(instance_name)
@@ -207,8 +243,14 @@ def get_bprop_mirror_micro_step_operator(self):
         assign.add_prim_attr("parameter_micro", 0)
     out_tensor = Tensor(1.0, mstype.float16)
     opt_shard = _get_enable_parallel_optimizer()
+    ln_print = P.Print()
+    reduce_sum = P.ReduceSum(keep_dims=False)
+    square = P.Square()
+    dump_local_norm = ms.get_auto_parallel_context("dump_local_norm")
 
     def bprop(x, z, out, dout):
+        if dump_local_norm:
+            z = F.depend(z, ln_print("dump local norm: ", param_name, reduce_sum(square((z)))))
         real_grad = z
         assign_out = dout
         if issubclass_(F.typeof(dout), mstype.tensor_type):
@@ -309,6 +351,9 @@ def get_bprop_micro_step_all_gather(self):
     """Generate bprop for _MicroStepAllGather"""
     fusion = self.get_attr_dict()["fusion"]
     mean_flag = self.get_attr_dict()["mean_flag"]
+    param_name = " "
+    if 'mirror_user_id' in self.get_attr_dict():
+        param_name = self.get_attr_dict()['mirror_user_id']
     do_mirror = False
     if self.group != "":
         do_mirror = self.get_attr_dict()["do_mirror"]
@@ -324,6 +369,10 @@ def get_bprop_micro_step_all_gather(self):
     dtype = P.DType()
     out_tensor = Tensor(1.0, mstype.float16)
     with_mirror_operator = self.get_attr_dict()["with_mirror_operator"]
+    ln_print = P.Print()
+    reduce_sum = P.ReduceSum(keep_dims=False)
+    square = P.Square()
+    dump_local_norm = ms.get_auto_parallel_context("dump_local_norm")
 
     def bprop(x, z, out, dout):
         if with_mirror_operator:
@@ -334,6 +383,8 @@ def get_bprop_micro_step_all_gather(self):
                 real_grad = F.tensor_mul(real_grad, scale)
             return (real_grad, cast(out_tensor, dtype(z)))
         z = F.depend(z, dout)
+        if dump_local_norm:
+            z = F.depend(z, ln_print("dump local norm: ", param_name, reduce_sum(square((z)))))
         if not do_mirror:
             return (z, cast(out_tensor, dtype(z)))
         real_grad = reduce_scatter(z)
@@ -529,16 +580,25 @@ def get_bprop_mirror_operator(self):
     group = self.get_attr_dict()['group']
     dev_num = self.get_attr_dict()['dev_num']
     mean_flag = self.get_attr_dict()['mean_flag']
+    param_name = " "
+    if 'mirror_user_id' in self.get_attr_dict():
+        param_name = self.get_attr_dict()['mirror_user_id']
+
     dev_num_r = 1.0
+    dump_local_norm = ms.get_auto_parallel_context("dump_local_norm")
     if dev_num > 1:
         dev_num_r = 1.0 / dev_num
         all_reduce = AllReduce(group=group)
         all_gather = AllGather(group=group)
         mul = P.Mul()
         cast = P.Cast()
+        ln_print = P.Print()
+        reduce_sum = P.ReduceSum(keep_dims=False)
+        square = P.Square()
 
         fusion = self.get_attr_dict()["fusion"]
         all_reduce.add_prim_attr("fusion", fusion)
+        parameter = " "
         if hasattr(self, 'parameter'):
             parameter = self.parameter
             all_reduce.add_prim_attr("parameter", parameter)
@@ -548,6 +608,9 @@ def get_bprop_mirror_operator(self):
             all_reduce.set_prim_instance_name(instance_name)
 
     def bprop(x, out, dout):
+        if dump_local_norm:
+            dout = F.depend(dout, ln_print("dump local norm: ", param_name, reduce_sum(square((dout)))))
+
         if dev_num == 1:
             return (dout,)
         if mean_flag:

mindspore/ops/_grad_experimental/grad_math_ops.py

@@ -18,12 +18,9 @@
 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 SilentCheck
-from mindspore.ops.operations._inner_ops import _MirrorSilentCheck
 from mindspore.ops.operations.math_ops import CumulativeLogsumexp
 from mindspore.ops.operations.math_ops import MatrixSolve
 from mindspore.ops.operations.math_ops import MatrixSolveLs
@@ -803,22 +800,3 @@ 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 dout.dtype == mstype.float16:
-            return (dout, out_tensor, out_tensor, out_tensor, out_tensor, out_tensor)
-        if loss_scale is not None:
-            gnorm = ops.norm(dout / loss_scale)
-        else:
-            gnorm = ops.norm(dout)
-        dx, _, _, _, _ = silent_check(gnorm, dout, pre_val, min_val, max_val, n_step)
-        return (dx, out_tensor, out_tensor, out_tensor, out_tensor, out_tensor)
-
-    return bporp

mindspore/ops/_vmap/vmap_array_ops.py

@@ -2113,6 +2113,7 @@ def get_split_vmap_rule(prim, axis_size):
 
     return vmap_rule
 
+
 @vmap_rules_getters.register(P.SearchSorted)
 def get_searchsorted_vmap_rule(prim, axis_size):
     """VmapRule for `SearchSorted`."""
@@ -2131,10 +2132,7 @@ def get_searchsorted_vmap_rule(prim, axis_size):
         if sorter is not None and sorter_dim is not None:
             sorter = _bdim_at_front(sorter, sorter_dim, axis_size)
 
-        dtype, _ = dtype_bdim
-        right, _ = right_bdim
-
-        outputs = prim(sequence, values, sorter, dtype, right)
+        outputs = prim(sequence, values, sorter, dtype_bdim[0], right_bdim[0])
 
         return outputs, 0
 

mindspore/ops/_vmap/vmap_math_ops.py

@@ -916,6 +916,23 @@ def get_isclose_vmap_rule(prim, axis_size):
 
     return vmap_rule
 
+
+@vmap_rules_getters.register(P.Round)
+def get_round_vmap_rule(prim, axis_size):
+    """VmapRule for round."""
+    if isinstance(prim, str):
+        prim = Primitive(prim)
+
+    def vmap_rule(x_bdim, decimal_bdim):
+        var, x_dim = x_bdim
+        decimal_var, decimal_dim = decimal_bdim
+        if decimal_dim is not None:
+            _raise_value_error("For vmap, the batch axis of decimal must be none.")
+        out = prim(var, decimal_var)
+        return out, x_dim
+
+    return vmap_rule
+
 get_assign_vmap_rule = vmap_rules_getters.register(P.AssignAdd)(get_assign_vmap_rule)
 get_assign_vmap_rule = vmap_rules_getters.register(P.AssignSub)(get_assign_vmap_rule)
 
@@ -949,7 +966,6 @@ get_unop_vmap_rule = vmap_rules_getters.register(P.Reciprocal)(get_unop_vmap_rul
 get_unop_vmap_rule = vmap_rules_getters.register(P.Inv)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Invert)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Rint)(get_unop_vmap_rule)
-get_unop_vmap_rule = vmap_rules_getters.register(P.Round)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Rsqrt)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register("Sigmoid")(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Sqrt)(get_unop_vmap_rule)

mindspore/ops/_vmap/vmap_nn_ops.py

@@ -517,7 +517,6 @@ def get_in_top_k_vmap_rule(prim, axis_size):
 
 @vmap_rules_getters.register(G.FastGeLUGrad)
 @vmap_rules_getters.register(G.HSwishGrad)
-@vmap_rules_getters.register(G.SoftShrinkGrad)
 def get_common_activation_grad_vmap_rule(prim, axis_size):
     """VmapRule for common activation grad operation."""
     prim_name = prim.name
@@ -547,6 +546,49 @@ def get_common_activation_grad_vmap_rule(prim, axis_size):
     return vmap_rule
 
 
+@vmap_rules_getters.register("SoftShrink")
+def get_softshrink_vmap_rule(prim, axis_size):
+    """VmapRule for `SoftShrink`."""
+    def vmap_rule(x_bdim, lambd_bdim):
+        var, dim = x_bdim
+        lambd, _ = lambd_bdim
+        out = prim(var, lambd)
+        return out, dim
+
+    return vmap_rule
+
+
+@vmap_rules_getters.register("SoftShrinkGrad")
+def get_softshrink_grad_vmap_rule(prim, axis_size):
+    """VmapRule for `SoftShrinkGrad`."""
+    prim_name = prim.name
+
+    def vmap_rule(dy_bdim, x_bdim, lambd_bdim):
+        x, x_dim = x_bdim
+        lambd, _ = lambd_bdim
+        dy, dy_dim = dy_bdim
+        x_shape = F.shape(x)
+        dy_shape = F.shape(dy)
+        if x_dim == dy_dim and x_shape == dy_shape:
+            out = prim(dy, x, lambd)
+            return out, x_dim
+
+        if F.rank(x):
+            x = _bdim_at_front(x, x_dim, 1)
+        if F.rank(dy):
+            dy = _bdim_at_front(dy, dy_dim, 1)
+        x_shape = F.shape(x)
+        dy_shape = F.shape(dy)
+        if x_shape != dy_shape:
+            raise RuntimeError("For {} vmap, input x shape is supposed to be the same as input dy shape "
+                               "after batch transforming, but got x_shape {}, dy_shape {}"
+                               .format(prim_name, x_shape, dy_shape))
+        out = prim(dy, x, lambd)
+        return out, 0
+
+    return vmap_rule
+
+
 @vmap_rules_getters.register("HShrink")
 def get_hshrink_vmap_rule(prim, axis_size):
     """VmapRule for `HShrink`."""
@@ -2196,7 +2238,6 @@ get_unop_vmap_rule = vmap_rules_getters.register(P.SeLU)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.HSigmoid)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Softplus)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.Softsign)(get_unop_vmap_rule)
-get_unop_vmap_rule = vmap_rules_getters.register(P.SoftShrink)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.GeLU)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.FastGeLU)(get_unop_vmap_rule)
 get_unop_vmap_rule = vmap_rules_getters.register(P.HSwish)(get_unop_vmap_rule)

mindspore/ops/auto_generate/cpp_create_prim_instance_helper.py

@@ -21,6 +21,7 @@ op_args_default_value = {
     "AdamW": {"amsgrad": False, "maximize": False},
     "AddExt": {"alpha": 1},
     "AddLayerNormV2": {"epsilon": 1e-5, "additionalOut": False},
+    "ApplyAdamW": {"max_grad_norm": None, "amsgrad": False, "maximize": False},
     "ApplyCamePart2": {"sum_r": None, "global_shape": None},
     "ApplyCamePart3": {"global_shape": None, "use_first_moment": False},
     "ApplyCamePart4": {"global_shape": None},
@@ -29,6 +30,7 @@ op_args_default_value = {
     "ArgMaxExt": {"dim": None, "keepdim": False},
     "Argmax": {"axis": -1, "output_type": mstype.int32},
     "ArgMaxWithValue": {"axis": 0, "keep_dims": False},
+    "ArgMinExt": {"dim": None, "keepdim": False},
     "Argmin": {"axis": -1, "output_type": mstype.int32},
     "ArgMinWithValue": {"axis": 0, "keep_dims": False},
     "AvgPool2DGrad": {"padding": 0, "ceil_mode": False, "count_include_pad": True, "divisor_override": None},
@@ -36,8 +38,8 @@ op_args_default_value = {
     "AvgPoolGrad": {"kernel_size": 1, "strides": 1, "pad_mode": 'VALID', "data_format": 'NCHW'},
     "AvgPool": {"kernel_size": 1, "strides": 1, "pad_mode": 'VALID', "data_format": 'NCHW'},
     "BatchMatMul": {"transpose_a": False, "transpose_b": False},
-    "BatchNormExt": {"training": False, "momentum": 0.1, "epsilon": 1e-5},
-    "BatchNormGradExt": {"training": False, "eps": 1e-5},
+    "BatchNormExt": {"running_mean": None, "runnning_var": None, "training": False, "momentum": 0.1, "epsilon": 1e-5},
+    "BatchNormGradExt": {"running_mean": None, "running_var": None, "saved_mean": None, "saved_rstd": None, "training": False, "eps": 1e-5},
     "BatchNormGradGrad": {"is_training": False, "epsilon": 1e-5, "data_format": 'NCHW'},
     "BatchNormGrad": {"is_training": False, "epsilon": 1e-5, "data_format": 'NCHW'},
     "BatchNormGradWithActivation": {"is_training": False, "epsilon": 1e-5, "data_format": 'NCHW'},
@@ -63,10 +65,13 @@ op_args_default_value = {
     "ConvolutionGrad": {"bias": None, "stride": 1, "padding": 0, "dilation": 1, "transposed": False, "output_padding": 0, "groups": 1, "output_mask": ()},
     "Convolution": {"bias": None, "stride": 1, "padding": 0, "dilation": 1, "transposed": False, "output_padding": 0, "groups": 1},
     "Correlate": {"mode": 'valid'},
+    "CountNonZero": {"dim": None},
+    "Cross": {"dim": -65530},
     "CumProd": {"exclusive": False, "reverse": False},
     "CumSum": {"exclusive": False, "reverse": False},
     "CumsumExt": {"dtype": None},
-    "DCT": {"axis": -1, "norm": 'BACKWARD', "forward": True, "grad": False},
+    "DCT": {"type": 2, "n": None, "axis": -1, "norm": None},
+    "DCTN": {"type": 2, "s": None, "axes": None, "norm": None},
     "Dense": {"bias": None},
     "Diagonal": {"offset": 0, "dim1": 0, "dim2": 1},
     "DivMod": {"rounding_mode": None},
@@ -75,13 +80,25 @@ op_args_default_value = {
     "EluExt": {"alpha": 1.0},
     "EluGradExt": {"alpha": 1.0},
     "Elu": {"alpha": 1.0},
+    "EmbeddingApplyAdaGrad": {"mask_zero": (0,), "padding_key": (0,), "padding_key_mask": (1,), "completion_key": (0,), "completion_key_mask": (1,), "_embedding_dim": 1, "_max_key_num": 1},
+    "EmbeddingApplyAdam": {"mask_zero": (0,), "padding_key": (0,), "padding_key_mask": (1,), "completion_key": (0,), "completion_key_mask": (1,), "_embedding_dim": 1, "_max_key_num": 1},
+    "EmbeddingApplyAdamW": {"ams_grad": (0,), "mask_zero": (0,), "padding_key": (0,), "padding_key_mask": (1,), "completion_key": (0,), "completion_key_mask": (1,), "_embedding_dim": 1, "_max_key_num": 1},
+    "EmbeddingApplyFtrl": {"mask_zero": (0,), "padding_key": (0,), "padding_key_mask": (1,), "completion_key": (0,), "completion_key_mask": (1,), "_embedding_dim": 1, "_max_key_num": 1},
+    "EmbeddingApplyRmsprop": {"mask_zero": (0,), "padding_key": (0,), "padding_key_mask": (1,), "completion_key": (0,), "completion_key_mask": (1,), "_embedding_dim": 1, "_max_key_num": 1},
+    "EmbeddingApplySgd": {"mask_zero": (0,), "padding_key": (0,), "padding_key_mask": (1,), "completion_key": (0,), "completion_key_mask": (1,), "_embedding_dim": 1, "_max_key_num": 1},
     "EmbeddingDenseBackward": {"padding_idx": None, "scale_grad_by_freq": False},
+    "EmbeddingFeatureMappingFileSize": {"only_offset_flag": True},
+    "EmbeddingFeatureMappingFind": {"num": 1},
+    "EmbeddingFeatureMappingImport": {"only_offset_flag": True, "num": 1},
     "Embedding": {"padding_idx": None, "max_norm": None, "norm_type": 2.0, "scale_grad_by_freq": False},
+    "EmbeddingTableEvict": {"steps_to_live": 0},
     "ExtractImagePatches": {"padding": 'VALID'},
     "FFNExt": {"expertTokens": None, "bias1": None, "bias2": None, "scale": None, "offset": None, "deqScale1": None, "deqScale2": None, "antiquant_scale1": None, "antiquant_scale2": None, "antiquant_offset1": None, "antiquant_offset2": None, "activation": 'fastgelu', "inner_precise": 0},
     "FFT2": {"s": None, "dim": (-2, -1), "norm": None},
     "FFT": {"n": None, "dim": -1, "norm": None},
+    "FFTOrtho": {"axes": None, "forward": True},
     "FFTWithSize": {"norm": 'backward', "onesided": True, "signal_sizes": ()},
+    "FFTFreq": {"d": 1.0, "dtype": None},
     "FFTN": {"s": None, "dim": None, "norm": None},
     "FFTShift": {"dim": None},
     "FillScalar": {"dtype": None},
@@ -90,23 +107,42 @@ op_args_default_value = {
     "FlashAttentionScore": {"real_shift": None, "drop_mask": None, "padding_mask": None, "attn_mask": None, "prefix": None, "actual_seq_qlen": None, "actual_seq_kvlen": None, "keep_prob": 1.0, "scale_value": 1.0, "pre_tokens": 2147483647, "next_tokens": 2147483647, "inner_precise": 0, "input_layout": 'BSH', "sparse_mode": 0},
     "FlattenExt": {"start_dim": 0, "end_dim": -1},
     "Gather": {"batch_dims": 0},
+    "GenerateEodMaskV2": {"start": 0, "steps": 1, "error_mode": 'cycle', "flip_mode": 'bitflip', "multiply_factor": 0.0, "bit_pos": 0, "flip_probability": 0.0},
     "GridSampler2DGrad": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
     "GridSampler2D": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
     "GridSampler3DGrad": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
     "GridSampler3D": {"interpolation_mode": 'bilinear', "padding_mode": 'zeros', "align_corners": False},
     "GroupNormGrad": {"dx_is_require": True, "dgamma_is_require": True, "dbeta_is_require": True},
     "GroupNorm": {"weight": None, "bias": None, "eps": 1e-5},
+    "HFFT2": {"s": None, "dim": (-2, -1), "norm": None},
+    "HFFT": {"n": None, "dim": -1, "norm": None},
+    "HFFTN": {"s": None, "dim": None, "norm": None},
+    "HistcExt": {"bins": 100, "min": 0, "max": 0},
     "HShrinkGrad": {"lambd": 0.5},
     "HShrink": {"lambd": 0.5},
+    "IDCT": {"type": 2, "n": None, "axis": -1, "norm": None},
+    "IDCTN": {"type": 2, "s": None, "axes": None, "norm": None},
     "IFFT2": {"s": None, "dim": (-2, -1), "norm": None},
     "IFFT": {"n": None, "dim": -1, "norm": None},
     "IFFTN": {"s": None, "dim": None, "norm": None},
     "IFFTShift": {"dim": None},
+    "IHFFT2": {"s": None, "dim": (-2, -1), "norm": None},
+    "IHFFT": {"n": None, "dim": -1, "norm": None},
+    "IHFFTN": {"s": None, "dim": None, "norm": None},
     "Im2ColExt": {"dilation": 1, "padding": 0, "stride": 1},
+    "IncreFlashAttention": {"attn_mask": None, "actual_seq_lengths": None, "pse_shift": None, "dequant_scale1": None, "quant_scale1": None, "dequant_scale2": None, "quant_scale2": None, "quant_offset2": None, "antiquant_scale": None, "antiquant_offset": None, "block_table": None, "kv_padding_size": None, "num_heads": 1, "input_layout": 'BSH', "scale_value": 1.0, "num_key_value_heads": 0, "block_size": 0, "inner_precise": 1},
     "IndexAddExt": {"alpha": 1},
-    "IRFFTGrad": {"n": None, "dim": -1, "norm": None},
+    "InplaceAddExt": {"alpha": 1},
+    "InplaceAddmm": {"beta": 1, "alpha": 1},
+    "InplaceAddsExt": {"alpha": 1},
+    "InsertGemV2InBackward": {"start": 0, "steps": 1, "error_mode": 'cycle', "flip_mode": 'bitflip', "multiply_factor": 0.0, "bit_pos": 0, "flip_probability": 0.0},
+    "IRFFT2": {"s": None, "dim": (-2, -1), "norm": None},
+    "IRFFTDouble": {"dim": -1},
     "IRFFT": {"n": None, "dim": -1, "norm": None},
+    "IRFFTN": {"s": None, "dim": None, "norm": None},
     "IsClose": {"rtol": 1e-05, "atol": 1e-08, "equal_nan": True},
+    "L1LossBackwardExt": {"reduction": 'mean'},
+    "L1LossExt": {"reduction": 'mean'},
     "LayerNormExt": {"weight": None, "bias": None, "eps": 1e-5},
     "LayerNormGradGrad": {"begin_norm_axis": 1, "begin_params_axis": 1},
     "LayerNormGrad": {"begin_norm_axis": 1, "begin_params_axis": 1},
@@ -116,10 +152,13 @@ op_args_default_value = {
     "LeakyReLUExt": {"negative_slope": 0.01},
     "LeakyReLUGradExt": {"negative_slope": 0.01, "is_result": False},
     "LinSpaceExt": {"dtype": None},
+    "LogSoftmaxExt": {"dim": None, "dtype": None},
     "LogSoftmaxGrad": {"axis": -1},
     "LogSoftmax": {"axis": -1},
     "LogitGrad": {"eps": -1.0},
     "Logit": {"eps": -1.0},
+    "LpNormV2": {"p": 2.0, "dim": None, "keepdim": False, "epsilon": 1e-12},
+    "LstsqV2": {"driver": None},
     "MatMul": {"transpose_a": False, "transpose_b": False},
     "MaxPoolGradWithIndices": {"strides": None, "pads": 0, "dilation": (1, 1), "ceil_mode": False, "argmax_type": mstype.int64},
     "MaxPoolGradWithMask": {"strides": None, "pads": 0, "dilation": (1, 1), "ceil_mode": False, "argmax_type": mstype.int64},
@@ -128,20 +167,30 @@ op_args_default_value = {
     "MaximumGradGrad": {"grad_x": True, "grad_y": True},
     "MaximumGrad": {"grad_x": True, "grad_y": True},
     "MeanExt": {"axis": None, "keep_dims": False, "dtype": None},
+    "MedianDim": {"dim": -1, "keepdim": False},
     "MinimumGrad": {"grad_x": True, "grad_y": True},
+    "MSELossExt": {"reduction": 'mean'},
+    "MSELossGradExt": {"reduction": 'mean'},
     "NanToNum": {"nan": None, "posinf": None, "neginf": None},
     "NLLLossGrad": {"reduction": 'mean', "ignore_index": -100},
     "NLLLoss": {"reduction": 'mean', "ignore_index": -100},
-    "Norm": {"ord": None, "dim": None, "keepdim": False, "dtype": None},
+    "Norm": {"p": 2.0, "dim": None, "keepdim": False, "dtype": None},
     "OneHotExt": {"axis": -1},
     "OneHot": {"axis": -1},
     "OnesLikeExt": {"dtype": None},
     "Ones": {"dtype": None},
+    "PagedAttentionMask": {"antiquant_scale": None, "antiquant_offset": None, "alibi_mask": None, "kv_cache_quant_mode": 'DEFAULT'},
+    "PagedAttention": {"antiquant_scale": None, "antiquant_offset": None, "attn_mask": None, "q_seq_lens": None, "kv_cache_quant_mode": 'DEFAULT'},
     "ProdExt": {"axis": None, "keep_dims": False, "dtype": None},
     "PromptKVCache": {"align_mode": 'LEFT'},
     "Qr": {"full_matrices": False},
     "RandExt": {"dtype": None},
     "RandLikeExt": {"dtype": None},
+    "RandIntLike": {"dtype": None},
+    "RandInt": {"dtype": None},
+    "RandnLike": {"dtype": None},
+    "Randn": {"dtype": None},
+    "RandpermExt": {"dtype": mstype.int64},
     "RandpermV2": {"seed": 0, "offset": 0, "dtype": mstype.int64},
     "Range": {"maxlen": 1000000},
     "ReduceAll": {"axis": None, "keep_dims": False},
@@ -165,16 +214,27 @@ op_args_default_value = {
     "ResizeNearestNeighbor": {"align_corners": False, "half_pixel_centers": False},
     "ResizeNearestNeighborV2Grad": {"align_corners": False, "half_pixel_centers": False},
     "ResizeNearestNeighborV2": {"align_corners": False, "half_pixel_centers": False},
-    "RFFTGrad": {"n": None, "dim": -1, "norm": None},
+    "RFFT2": {"s": None, "dim": (-2, -1), "norm": None},
     "RFFT": {"n": None, "dim": -1, "norm": None},
+    "RFFTFreq": {"d": 1.0, "dtype": None},
+    "RFFTN": {"s": None, "dim": None, "norm": None},
     "RmsNorm": {"epsilon": 1e-6},
+    "Roll": {"axis": None},
+    "RotaryPositionEmbeddingGrad": {"dx": None, "mode": 0},
+    "RotaryPositionEmbedding": {"mode": 0},
+    "Round": {"decimals": 0},
     "ScalarToTensor": {"dtype": None},
+    "Scatter": {"reduce": 'none'},
+    "ScatterValue": {"reduce": 'none'},
     "SearchSorted": {"sorter": None, "dtype": mstype.int64, "right": False},
     "SequenceConcat": {"axis": 0},
+    "SilentCheckV2": {"c_min_steps": 7, "c_thresh_l1": 1000000.0, "c_coeff_l1": 100000.0, "c_thresh_l2": 10000.0, "c_coeff_l2": 5000.0, "npu_asd_detect": 1},
     "SoftmaxBackward": {"dim": -1},
     "Softmax": {"axis": -1},
     "SoftplusExt": {"beta": 1, "threshold": 20},
     "SoftplusGradExt": {"beta": 1, "threshold": 20},
+    "SoftShrinkGrad": {"lambd": 0.5},
+    "SoftShrink": {"lambd": 0.5},
     "SolveTriangular": {"trans": 0, "lower": False, "unit_diagonal": False},
     "SortExt": {"dim": -1, "descending": False, "stable": False},
     "Split": {"axis": 0, "output_num": 1},
@@ -184,11 +244,20 @@ op_args_default_value = {
     "StridedSlice": {"begin_mask": 0, "end_mask": 0, "ellipsis_mask": 0, "new_axis_mask": 0, "shrink_axis_mask": 0},
     "SubExt": {"alpha": 1},
     "SumExt": {"dim": None, "keepdim": False, "dtype": None},
+    "SwigluGrad": {"dim": -1},
+    "Swiglu": {"dim": -1},
+    "TensorScatterElements": {"axis": 0, "reduce": 'none'},
     "TopkExt": {"dim": -1, "largest": True, "sorted": True},
+    "TopKRouter": {"drop_type": 0},
+    "TraceV2Grad": {"offset": 0, "axis1": 1, "axis2": 0},
+    "TraceV2": {"offset": 0, "axis1": 1, "axis2": 0, "dtype": None},
+    "TrilExt": {"diagonal": 0},
     "Triu": {"diagonal": 0},
     "TupleToTensor": {"dtype": None},
     "Unique2": {"sorted": True, "return_inverse": False, "return_counts": False},
     "UnstackExt": {"axis": 0},
+    "UpsampleBicubic2DGrad": {"output_size": None, "scales": None, "align_corners": False},
+    "UpsampleBicubic2D": {"output_size": None, "scales": None, "align_corners": False},
     "UpsampleBilinear2DGrad": {"output_size": None, "scales": None, "align_corners": False},
     "UpsampleBilinear2D": {"output_size": None, "scales": None, "align_corners": False},
     "UpsampleLinear1DGrad": {"output_size": None, "scales": None, "align_corners": False},
@@ -203,12 +272,13 @@ op_args_default_value = {
     "UpsampleTrilinear3D": {"output_size": None, "scales": None, "align_corners": False},
     "ZerosLikeExt": {"dtype": None},
     "Zeros": {"dtype": None},
+    "AddRmsNormQuantV2": {"epsilon": 1e-5},
     "DynamicQuantExt": {"smooth_scales": None},
     "FusedInferAttentionScore": {"pse_shift": None, "attn_mask": None, "actual_seq_lengths": None, "actual_seq_lengths_kv": None, "dequant_scale1": None, "quant_scale1": None, "dequant_scale2": None, "quant_scale2": None, "quant_offset2": None, "antiquant_scale": None, "antiquant_offset": None, "block_table": None, "query_padding_size": None, "kv_padding_size": None, "scale_value": 1.0, "pre_tokens": 2147483647, "next_tokens": 2147483647, "input_layout": 'BSH', "num_key_value_heads": 0, "sparse_mode": 0, "inner_precise": 1, "block_size": 0, "antiquant_mode": 0, "softmax_lse_flag": False},
     "GroupedMatmul": {"bias": None, "scale": None, "offset": None, "antiquant_scale": None, "antiquant_offset": None, "group_list": None, "split_item": 0, "group_type": -1},
     "KVCacheScatterUpdate": {"reduce": 'none'},
     "MoeFinalizeRouting": {"x2": None, "bias": None, "scales": None, "expanded_row_idx": None, "expanded_expert_idx": None},
-    "QuantBatchMatmul": {"offset": None, "bias": None, "transpose_x1": False, "transpose_x2": False, "dtype": mstype.float16},
+    "QuantBatchMatmul": {"offset": None, "bias": None, "pertokenScaleOptional": None, "transpose_x1": False, "transpose_x2": False, "dtype": mstype.float16},
     "QuantV2": {"sqrt_mode": False, "rounding_mode": 'ROUND', "dst_type": mstype.int8},
     "WeightQuantBatchMatmul": {"antiquant_offset": None, "quant_scale": None, "quant_offset": None, "bias": None, "transpose_x": False, "transpose_weight": False, "antiquant_group_size": 0},
 }
@@ -216,16 +286,30 @@ op_args_default_value = {
 op_labels = {
     "AdamWeightDecay": {"side_effect_mem": True},
     "AdamW": {"side_effect_mem": True},
+    "ApplyAdamW": {"side_effect_mem": True},
     "AssignAdd": {"side_effect_mem": True},
     "Assign": {"side_effect_mem": True},
+    "CopyExt": {"side_effect_mem": True},
     "DecoderKVCache": {"side_effect_mem": True},
     "DropoutExt": {"side_effect_hidden": True},
     "DropoutGenMaskExt": {"side_effect_hidden": True},
     "Dropout": {"side_effect_hidden": True},
+    "EmbeddingApplyAdaGrad": {"_process_node_engine_id": 'PS'},
+    "EmbeddingApplyAdam": {"_process_node_engine_id": 'PS'},
+    "EmbeddingApplyAdamW": {"_process_node_engine_id": 'PS'},
+    "EmbeddingApplyFtrl": {"_process_node_engine_id": 'PS'},
+    "EmbeddingApplyRmsprop": {"_process_node_engine_id": 'PS'},
+    "EmbeddingApplySgd": {"_process_node_engine_id": 'PS'},
     "Embedding": {"side_effect_mem": True},
+    "EmbeddingTableEvict": {"_process_node_engine_id": 'PS'},
     "Generator": {"side_effect_mem": True},
+    "InplaceAddExt": {"side_effect_mem": True},
+    "InplaceAddmm": {"side_effect_mem": True},
+    "InplaceAddsExt": {"side_effect_mem": True},
     "Log": {"cust_aicpu": 'Log', "base": -1.0, "scale": 1.0, "shift": 0.0},
     "PromptKVCache": {"side_effect_mem": True},
     "ReshapeAndCache": {"side_effect_mem": True},
     "ResizeD": {"mode": 'linear'},
+    "SilentCheckV2": {"side_effect_mem": True},
+    "KVCacheScatterUpdate": {"side_effect_mem": True},
 }