mindspore 2.5.0__cp311-cp311-win_amd64.whl → 2.6.0__cp311-cp311-win_amd64.whl

mindspore/ops/composite/multitype_ops/ones_like_impl.py

@@ -29,13 +29,17 @@ ones_like_leaf = base.MultitypeFuncGraph('ones_like', True)
 using ".register" decorator.
 """
 
+_ones_like_leaf_for_grad = base.MultitypeFuncGraph('_ones_like_for_grad', True)
 
+
+@_ones_like_leaf_for_grad.register("TypeType")
 @ones_like_leaf.register("TypeType")
 def _ones_like_type_type(x):
     """Returns x because x is a type. This is usually used in backprop progress."""
     return x
 
 
+@_ones_like_leaf_for_grad.register("Number")
 @ones_like_leaf.register("Number")
 def _ones_like_scalar(x):
     """Returns 1 which has the same dtype as x where x is a scalar."""
@@ -43,12 +47,14 @@ def _ones_like_scalar(x):
     return F.scalar_cast(1.0, t)
 
 
+@_ones_like_leaf_for_grad.register("Tensor")
 @ones_like_leaf.register("Tensor")
 def _ones_like_tensor(x):
     """Returns a tensor with the same shape and dtype as x and all elements are 1."""
     return P.OnesLike()(x)
 
 
+@_ones_like_leaf_for_grad.register("COOTensor")
 @ones_like_leaf.register("COOTensor")
 def _ones_like_coo_tensor(x):
     """Returns a tensor with the same shape and dtype as x and all elements are 1."""
@@ -57,12 +63,14 @@ def _ones_like_coo_tensor(x):
     return F.make_coo_tensor(F.coo_tensor_get_indices(x), values, F.coo_tensor_get_dense_shape(x))
 
 
+@_ones_like_leaf_for_grad.register("CSRTensor")
 @ones_like_leaf.register("CSRTensor")
 def _ones_like_csr_tensor(x):
     """Returns a tensor with the same shape and dtype as x and all elements are 1."""
     return F.make_csr_tensor(x.indptr, x.indices, ones_like(x.values), x.shape)
 
 
+@_ones_like_leaf_for_grad.register("Function")
 @ones_like_leaf.register("Function")
 def _ones_like_func(x):
     """
@@ -78,12 +86,20 @@ def _ones_like_func(x):
     return F.environ_create()
 
 
+@_ones_like_leaf_for_grad.register("None")
 @ones_like_leaf.register("None")
 def _ones_like_none(x):
     """Returns none"""
     return None
 
 
+# pylint: disable=protected-access
+@_ones_like_leaf_for_grad._register_default(False)
+def default_ones_like_leaf_for_grad(x):
+    """Default function for _ones_like_leaf_for_grad."""
+    return x
+
+
 ones_like = base.HyperMap(ones_like_leaf)
 """
 `ones_like` is a function which can generate a graph of `ones_like` operation according to input tensor dtype.
@@ -94,3 +110,5 @@ Example:
     >>> input = Tensor([2, 3], mindspore.float16)
     >>> ones = ones_like(input) # The dtype of ones is mindspore.float16
 """
+
+_ones_like_for_grad = base.HyperMap(_ones_like_leaf_for_grad)

mindspore/ops/composite/multitype_ops/pow_impl.py

@@ -52,7 +52,8 @@ def _scalar_pow_tensor(x, y):
     return F.tensor_pow(x, y)
 
 
-@pow_.register_default()
+# pylint: disable=protected-access
+@pow_._register_default()
 def default_pow(x, y):
     """Default function for pow."""
     return x ** y

mindspore/ops/composite/multitype_ops/right_shift_impl.py

@@ -31,7 +31,8 @@ def _right_shift_scalar(x, y):
     return inner.bit_right_shift(x, y)
 
 
-@right_shift.register_default()
+# pylint: disable=protected-access
+@right_shift._register_default()
 def default_right_shift(x, y):
     """Default function for right shift."""
     return x >> y

mindspore/ops/composite/multitype_ops/setitem_impl.py

@@ -877,7 +877,8 @@ def _map_tensor_setitem(map_tensor, key_tensor, value_tensor):
     return map_tensor
 
 
-@setitem.register_default()
+# pylint: disable=protected-access
+@setitem._register_default()
 def default_getitem(x, y, z):
     """Default function for setitem."""
     x[y] = z

mindspore/ops/composite/multitype_ops/sub_impl.py

@@ -110,7 +110,8 @@ def _cootensor_sub_tensor(x, y):
     return F.tensor_scatter_add(-y, x.indices, x.values)
 
 
-@sub.register_default()
+# pylint: disable=protected-access
+@sub._register_default()
 def default_sub(x, y):
     """Default function for sub."""
     return x - y

mindspore/ops/function/__init__.py

@@ -27,7 +27,16 @@ from . import (
     linalg_func,
     clip_func,
     fft_func,
-    reshard_func
+    reshard_func,
+    sparse_func,
+    random_func,
+    debug_func,
+    image_func,
+    spectral_func,
+    vmap_func,
+    sparse_unary_func,
+    other_func,
+    _add_attr_func
 )
 from .array_func import (
     unique,
@@ -70,6 +79,7 @@ from .array_func import (
     select_scatter,
     scalar_to_array,
     scalar_to_tensor,
+    take,
     tuple_to_array,
     expand_dims,
     squeeze,
@@ -217,6 +227,7 @@ from .math_func import (
     sqrt,
     square,
     tensor_mul,
+    tensor_muls,
     mul,
     multiply,
     digamma,
@@ -241,6 +252,7 @@ from .math_func import (
     lcm,
     tensor_exp,
     einsum,
+    einsum_ext,
     view_as_real,
     var,
     var_mean,
@@ -448,6 +460,7 @@ from .math_func import (
     count_nonzero,
     tensor_dot,
     vecdot,
+    multi_scale_deformable_attn_function,
     dot,
     isnan_ext,
     batch_dot,
@@ -466,6 +479,8 @@ from .nn_func import (
     avg_pool3d,
     max_pool2d,
     max_pool3d,
+    moe_token_permute,
+    moe_token_unpermute,
     batch_norm,
     add_rms_norm,
     rms_norm,
@@ -487,6 +502,7 @@ from .nn_func import (
     flipud,
     fractional_max_pool2d,
     fractional_max_pool3d,
+    speed_fusion_attention,
     pixel_shuffle,
     pixel_unshuffle,
     hardshrink,
@@ -495,6 +511,7 @@ from .nn_func import (
     incre_flash_attention,
     prompt_flash_attention,
     flash_attention_score,
+    fused_infer_attention_score,
     intopk,
     interpolate,
     upsample,
@@ -514,7 +531,6 @@ from .nn_func import (
     softsign,
     silu,
     selu,
-    soft_margin_loss,
     softmax,
     softmin,
     softshrink,
@@ -539,6 +555,7 @@ from .nn_func import (
     conv2d,
     conv_transpose2d,
     sigmoid,
+    soft_margin_loss,
     logsigmoid,
     relu,
     relu_,
@@ -645,6 +662,7 @@ from .grad import (
     vjp,
     linearize,
     stop_gradient,
+    stop_gradient_,
     get_grad
 )
 from .debug_func import (
@@ -766,7 +784,14 @@ from .other_func import (
     depend,
     partial,
     rotary_position_embedding,
+    move_to,
+    moe_init_routing_v2
 )
+
+from ._add_attr_func import (
+    _add_attr,
+)
+
 from .reshard_func import (
     reshard,
 )
@@ -790,5 +815,6 @@ __all__.extend(sparse_unary_func.__all__)
 __all__.extend(clip_func.__all__)
 __all__.extend(fft_func.__all__)
 __all__.extend(other_func.__all__)
+__all__.extend(_add_attr_func.__all__)
 __all__.extend(reshard_func.__all__)
 __all__.sort()

mindspore/ops/function/_add_attr_func.py

@@ -0,0 +1,58 @@
+# Copyright 2025 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.
+# ============================================================================
+"""_add_attr"""
+
+from typing import Callable
+import mindspore as ms
+from mindspore._c_expression import AddAttr_
+from mindspore.ops.primitive import Primitive
+
+
+class AddAttr(AddAttr_):
+    """
+    Config primitive attributes for operators in function.
+    """
+    def __init__(self):
+        super().__init__("AddAttr")
+        self.addattr_fn = None
+        self.fn = None
+        self.attrs_dict = None
+
+    def __call__(self, fn, **kwargs):
+        if not isinstance(kwargs, dict):
+            raise TypeError(f"the parameter 'kwargs' must be dict type, but got:{type(kwargs)}")
+        if (not isinstance(fn, Callable)) or isinstance(fn, Primitive):
+            raise TypeError(f"the parameter 'fn' must be callable type except Primitive, but got:{type(fn)}")
+
+        if self.addattr_fn is not None and self.addattr_fn == fn and \
+           self.fn is not None and self.fn == fn and \
+           self.attrs_dict is not None and self.attrs_dict == kwargs:
+            return self.addattr_fn
+
+        add_attr_ = AddAttr()
+        attr_kv_pair = tuple(kwargs.items())
+        @ms.common.jit
+        def add_attr_fn(*args):
+            return add_attr_(fn, attr_kv_pair)(*args)
+
+        self.addattr_fn = add_attr_fn
+        self.fn = fn
+        self.attrs_dict = kwargs
+        return self.addattr_fn
+
+def _add_attr(fn, **kwargs):
+    return  AddAttr()(fn, **kwargs)
+
+__all__ = ["_add_attr"]