mindspore 2.6.0__cp39-none-any.whl → 2.7.0rc1__cp39-none-any.whl

mindspore/ops/operations/debug_ops.py

@@ -30,7 +30,6 @@ from mindspore.ops.primitive import prim_attr_register, Primitive, PrimitiveWith
 from mindspore._checkparam import check_hook_fn
 from mindspore.ops import operations as P
 
-
 SUMMARY_TENSOR_CACHE = []
 
 
@@ -463,15 +462,67 @@ class InsertGradientOf(Primitive):
         self.f = f
 
 
+class DumpGradient(Primitive):
+    """
+        The `DumpGradient` Primitive is a hook, used to dump dout which pass to `x`.
+
+        Inputs:
+            - **path** (str) - The path of the file to be saved.
+            - **x** (Tensor) - Input Tensor of any dimension.
+            - **input_output** (str) - support value should be one of ['in', 'out'].
+
+        Supported Platforms:
+            ``Ascend``
+
+        Examples:
+            >>> import numpy as np
+            >>> import mindspore as ms
+            >>> from mindspore import ops
+            >>> from mindspore import Tensor
+            >>> ms.set_context(mode=ms.PYNATIVE_MODE)
+            >>> ms.set_device(device_target="Ascend")
+            >>> dg = ops.DumpGradient()
+            >>> def dout_dump_test(x, y):
+            ...     x = dg("x_dout.npy", x, 'out')
+            ...     print(f"x value is {x}")
+            ...     z = x * y
+            ...     return z
+            >>> ms_grad = ms.grad(dout_dump_test, grad_position=(0,1))
+            >>> x_grad, y_grad = ms_grad(Tensor(1, ms.float32), Tensor(2, ms.float32))
+            >>> print(f"x grad is {x_grad}, y_grad is {y_grad}")
+            >>> x_grad_npy = np.load("x_dout.npy")
+            >>> print(f"load x_grad from npy, x_grad is {x_grad_npy}")
+            x value is 1.0
+            x grad is 2.0, y grad is 1.0
+            load x_grad from npy, x_grad is array(2., dtype=float32)
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        pass
+
+    def _dump_hook(self, dout):
+        P.TensorDump()(self.bwd_dump_path, dout)
+        return dout
+
+    def __call__(self, path, x, input_output):
+        self.bwd_dump_path = path
+        x = P.InsertGradientOf(self._dump_hook)(x)
+        return x
+
+
 class Morph(PrimitiveWithInfer):
     """
     The `Morph` Primitive is used to encapsulate a user-defined function `fn`, allowing it to be used as a custom
     Primitive.
-    The primary application scenario of the `Morph` Primitive is in the auto-parallel case after `GRAPH_MODE` mode,
-    where collective communication operators are used within the user-defined `fn` to implement custom parallel
-    computation logic, especially in scenarios where `fn` involves dynamic shapes.
+
+    The `Morph` Primitive is primarily designed for custom graph optimization in GRAPH mode. For example, it supports
+    encapsulation of irregular collective communications (such as :func:`mindspore.ops.AlltoAllV`) in distributed
+    auto-parallel training scenarios.
+
     When the `Morph` Primitive is applied to inputs, it is actually the encapsulated user-defined function `fn` that is
     applied to the inputs.
+
     The main difference between the `Morph` Primitive and :func:`mindspore.ops.Custom` is that the former is expanded
     and replaced by the user-defined `fn` before automatic differentiation, so there is no need to implement a backward
     function.

mindspore/ops/operations/image_ops.py

@@ -102,19 +102,19 @@ class AdjustContrastv2(Primitive):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
-    >>> images = Tensor([[[1.0, 2.0, 3.0],
-    ...       [4.0, 5.0, 6.0]],
-    ...     [[7.0, 8.0, 9.0],
-    ...       [10.0, 11.0, 12.0]]], mstype.float32)
-    >>> contrast_factor = Tensor(2., mstype.float32)
-    >>> adjustcontrastv2 = AdjustContrastv2()
-    >>> output = adjustcontrastv2(images, contrast_factor)
-    >>> print(output)
-    [[[-3.5 -2.5 -1.5]
-      [ 2.5  3.5  4.5]]
-    <BLANKLINE>
-     [[ 8.5  9.5 10.5]
-      [14.5 15.5 16.5]]]
+        >>> images = Tensor([[[1.0, 2.0, 3.0],
+        ...       [4.0, 5.0, 6.0]],
+        ...     [[7.0, 8.0, 9.0],
+        ...       [10.0, 11.0, 12.0]]], mstype.float32)
+        >>> contrast_factor = Tensor(2., mstype.float32)
+        >>> adjustcontrastv2 = AdjustContrastv2()
+        >>> output = adjustcontrastv2(images, contrast_factor)
+        >>> print(output)
+        [[[-3.5 -2.5 -1.5]
+          [ 2.5  3.5  4.5]]
+        <BLANKLINE>
+         [[ 8.5  9.5 10.5]
+          [14.5 15.5 16.5]]]
     """
 
     @prim_attr_register

mindspore/ops/operations/manually_defined/ops_def.py

@@ -26,7 +26,6 @@ from mindspore.ops._primitive_cache import _get_cache_prim
 from mindspore.ops._utils import arg_handler as handler
 from mindspore.ops._utils.arg_dtype_cast import DtypeToEnum
 from mindspore.common import Tensor, CSRTensor, COOTensor
-from mindspore.common._stub_tensor import _convert_stub
 from mindspore._c_expression import typing
 from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._c_expression import pyboost_cast, pyboost_tile, pyboost_zeros, pyboost_ones, pyboost_type_as
@@ -1057,8 +1056,8 @@ class Tile(Primitive):
     def __call__(self, input, dims):
         # Add for jit context.
         if jit_context() and jit_context().compiled:
-            return None
-        res = _convert_stub(pyboost_tile(self, [input, dims]))
+            return jit_context().default_output()
+        res = pyboost_tile(self, [input, dims])
         # Add for jit context.
         if jit_context():
             if validator.is_stub_tensor(res):
@@ -1066,7 +1065,6 @@ class Tile(Primitive):
             return jit_context().run_op(self, res, input, dims)
         return res
 
-    # pylint: disable=missing-docstring
     def check_elim(self, *args):
         base_tensor, dims = args
         if not isinstance(base_tensor, Tensor):
@@ -1230,11 +1228,11 @@ class Cast(Primitive):
     def __call__(self, input_x, dtype):
         # Add for jit context.
         if jit_context() and jit_context().compiled:
-            return None
+            return jit_context().default_output()
         should_elim, output = self.check_elim(input_x, dtype)
         if should_elim:
             return output
-        res = _convert_stub(pyboost_cast(self, [input_x, dtype_to_type_id('Cast', 'dtype', dtype)]))
+        res = pyboost_cast(self, [input_x, dtype_to_type_id('Cast', 'dtype', dtype)])
         # Add for jit context.
         if jit_context():
             if validator.is_stub_tensor(res):
@@ -1293,7 +1291,7 @@ class TypeAs(Primitive):
     def __call__(self, input, other):
         if input.dtype == other.dtype:
             return input
-        return _convert_stub(pyboost_type_as(self, [input, other]))
+        return pyboost_type_as(self, [input, other])
 
 
 def to_sequence(val):
@@ -2070,9 +2068,9 @@ class Ones(Primitive):
     def __call__(self, size, type=None):
         # Add for jit context.
         if jit_context() and jit_context().compiled:
-            return None
-        res = _convert_stub(pyboost_ones(self, [size, type if type is None \
-            else handler.dtype_to_type_id('Ones', 'type', type)]))
+            return jit_context().default_output()
+        res = pyboost_ones(self, [size, type if type is None \
+            else handler.dtype_to_type_id('Ones', 'type', type)])
         # Add for jit context.
         if jit_context():
             if validator.is_stub_tensor(res):
@@ -2130,9 +2128,9 @@ class Zeros(Primitive):
     def __call__(self, size, type=None):
         # Add for jit context.
         if jit_context() and jit_context().compiled:
-            return None
-        res = _convert_stub(pyboost_zeros(self, [size, type if type is None else \
-            handler.dtype_to_type_id('Zeros', 'type', type)]))
+            return jit_context().default_output()
+        res = pyboost_zeros(self, [size, type if type is None else \
+            handler.dtype_to_type_id('Zeros', 'type', type)])
         # Add for jit context.
         if jit_context():
             if validator.is_stub_tensor(res):
@@ -2217,7 +2215,8 @@ def flash_attention_score(query, key, value, head_num, real_shift=None, drop_mas
         keep_prob (double, optional): The keep probability of dropout. Value range is (0.0, 1.0]. When `keep_prob`
             is 1.0, `drop_mask` should be None.
             Default: ``1.0``.
-        scalar_value (double, optional): The scale factor of score. Generally, the value is 1.0 / (D ** 0.5).
+        scalar_value (double, optional): The scale value indicating the scale coefficient, which is used as the
+            scalar of Muls in the calculation. Generally, the value is 1.0 / (D ** 0.5).
             Default: ``1.0``.
         pre_tokens (int, optional): Parameter for sparse computation, represents how many tokens are counted forward.
             When `sparse_mode` is set to 1, 2, 3, or 5, this parameter does not take effect.
@@ -2599,8 +2598,8 @@ def fused_infer_attention_score(query, key, value, *, pse_shift=None, atten_mask
         taking exp, and then the sum is computed to obtain softmax_sum. Finally, the log of softmax_sum is taken,
         and softmax_max is added to obtain softmax_lse. The softmax_lse is only calculated when softmax_lse_flag
         is True, and the shape would be :math:`(B, N, Q\_S, 1)`. If softmax_lse_flag is False, then a tensor with
-        shape :math:`(1)` filled with zeros would be returned. In graph mode with JitConfig set to O2, please ensure
-        that the softmax_lse_flag is enabled before using softmax_lse; otherwise, an exception will occur.
+        shape :math:`(1)` filled with zeros would be returned. In GE backend, please ensure that the softmax_lse_flag
+        is enabled before using softmax_lse; otherwise, an exception will occur.
 
     Constraints:
         - Full Inference Scenario (Q_S > 1):

mindspore/ops/operations/math_ops.py

@@ -1231,7 +1231,7 @@ class Heaviside(Primitive):
             0, & \text { if x }<0 \\
             \text { values, } & \text { if x }==0 \\
             1, & \text { if x }>0
-            \end{array}\right
+            \end{array}\right.
 
     .. warning::
         This is an experimental API that is subject to change or deletion.
@@ -3354,9 +3354,8 @@ class ComplexAbs(Primitive):
 
     Examples:
         >>> import mindspore
-        >>> import numpy as np
         >>> from mindspore import Tensor, ops
-        >>> x = Tensor(np.asarray(np.complex(3+4j)), mindspore.complex64)
+        >>> x = Tensor(3+4j, mindspore.complex64)
         >>> complex_abs = ops.ComplexAbs()
         >>> output = complex_abs(x)
         >>> print(output)
@@ -3630,7 +3629,7 @@ class MatrixSolveLs(Primitive):
         TypeError: If `l2_regularizer` is not float64.
         TypeError: If `fast` is not bool.
         ValueError: If dimensions of `matrix` or `rhs` is less than 2.
-        ValueError: If shape of `matrix` dose not match the shape of `rhs`.
+        ValueError: If shape of `matrix` does not match the shape of `rhs`.
 
     Supported Platforms:
         ``CPU``

mindspore/ops/operations/nn_ops.py

@@ -32,7 +32,6 @@ from mindspore.ops.primitive import PrimitiveWithCheck
 from mindspore.ops.primitive import prim_attr_register
 from mindspore.run_check._check_version import AscendEnvChecker
 from mindspore._c_expression import pyboost_all_finite
-from mindspore.common._stub_tensor import _convert_stub
 from ..auto_generate import (CeLU, Flatten, LogSoftmax, LogSoftmaxExt, GLU, ReLU, ReLU6, Dense, Tanh,
                              Elu, Sigmoid, Softmax, SoftplusExt, HSwish, HSigmoid, AvgPool, BiasAdd,
                              NLLLoss, OneHot, GeLU, FastGeLU, PReLU, RmsNorm, IncreFlashAttention, MSELossExt,
@@ -7120,8 +7119,8 @@ class Conv3DTranspose(Primitive):
         self.format = validator.check_string(data_format, ['NCDHW'], 'format', self.name)
         self.add_prim_attr('data_format', self.format)
 
-        self.output_padding = _check_3d_int_or_tuple('output_padding', output_padding, self.name,
-                                                     allow_five=False, ret_five=True, greater_zero=False)
+        self.output_padding = _check_3d_int_or_tuple('output_padding', output_padding, self.name, allow_five=False,
+                                                     ret_five=True, greater_zero=False, pad_value=0)
         output_padding_ = (self.output_padding[2], self.output_padding[3], self.output_padding[4])
         if self.pad_mode != 'pad' and output_padding_ != (0, 0, 0):
             raise ValueError(f"For '{self.name}', the 'output_padding' must be zero or (0, 0, 0) "
@@ -7371,8 +7370,8 @@ class ApplyAdagradDA(Primitive):
         >>> global_step = Tensor(2, mstype.int32)
         >>> output = net(grad, lr, l1, l2, global_step)
         >>> print(output)
-        [[-0.00073906, -0.00136889],
-         [-0.00059699, -0.00142478]]
+        [[-0.00073906 -0.00136889]
+         [-0.00059699 -0.00142478]]
     """
 
     __mindspore_signature__ = (
@@ -9258,4 +9257,4 @@ class AllFinite(Primitive):
                     "in the current environment does not support AllFinite.")
 
     def __call__(self, *args):
-        return _convert_stub(pyboost_all_finite(self, args))
+        return pyboost_all_finite(self, args)

mindspore/ops/primitive.py

@@ -170,10 +170,13 @@ class Primitive(Primitive_):
                 raise TypeError(f'The element of strategy must be tuple/Layout type, but got:{type(in_ele)}')
             if isinstance(in_ele, tuple):
                 for in_value in in_ele:
-                    if not isinstance(in_value, int) and self.name not in SUPPORTED_TUPLE_IN_TUPLE_STRATEGY:
+                    if not isinstance(in_value, int) and self.name not in SUPPORTED_TUPLE_IN_TUPLE_STRATEGY \
+                            and not self.attrs.get("self_define_shard", False):
                         raise TypeError(f'The {log_info}: {strategy} of {self.name} is not valid,'
                                         f' the value of strategy must be int type, but got:{type(in_value)}')
-                    if isinstance(in_value, Layout) and (self.name in SUPPORTED_TUPLE_IN_TUPLE_STRATEGY):
+                    if isinstance(in_value, Layout) and (
+                            self.name in SUPPORTED_TUPLE_IN_TUPLE_STRATEGY or self.attrs.get("self_define_shard",
+                                                                                             False)):
                         is_layout.append(True)
                         continue
                 is_layout.append(False)
@@ -315,7 +318,7 @@ class Primitive(Primitive_):
             out_is_layout = self._check_shard_strategy(out_strategy, "out_strategy")
         is_layout = in_is_layout if in_is_layout is not None else out_is_layout
         if out_is_layout is not None and is_layout != out_is_layout and \
-                self.name not in SUPPORTED_TUPLE_IN_TUPLE_STRATEGY:
+                self.name not in SUPPORTED_TUPLE_IN_TUPLE_STRATEGY and not self.attrs.get("self_define_shard", False):
             raise ValueError(f'The in_strategy type must equal to the out_strategy type, '
                              f'one using tuple(tuple) and the other using tuple(Layout) is not allowed.')
 
@@ -409,12 +412,6 @@ class Primitive(Primitive_):
             return output
         return _run_op(self, self.name, args)
 
-    def __getstate__(self):
-        return self.__dict__
-
-    def __setstate__(self, d):
-        self.__dict__.update(d)
-
     def __deepcopy__(self, memo):
         return type(self)(**self.init_attrs)
 
@@ -1032,7 +1029,6 @@ def _run_op(obj, op_name, args):
     res = _pynative_executor.run_op_async(obj, op_name, args)
     # Add for jit context.
     if jit_context():
-        # todo support TensorPy
         return jit_context().run_op(obj, res, *args)
     return res
 

mindspore/ops/tensor_method.py

@@ -21,7 +21,7 @@ from mindspore.ops import operations as P
 from mindspore.ops import functional as F
 from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite.multitype_ops._compile_utils import (
-    sequence_to_tensor, _tensor_sub, _tensor_pow, _tensor_div, _tensor_floordiv
+    sequence_to_tensor, _tensor_sub, _tensor_pow, _tensor_div, _tensor_floordiv, _tensor_mod
 )
 from mindspore.ops.auto_generate.gen_ops_prim import (
     inplace_scatter_src_op, inplace_scatter_src_reduce_op, inplace_scatter_value_op, inplace_scatter_value_reduce_op,
@@ -503,16 +503,16 @@ def tensor_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
 
 
 # 8 any
-def tensor_any(x, axis=None, keep_dims=False):
+def reduce_tensor_any(x, axis=None, keep_dims=False):
     if axis is None:
         axis = ()
     return any(x, axis, keep_dims)
 
 
-def deprecated_tensor_any(x, dim=None, keepdim=False):
+def tensor_any(input, dim=None, keepdim=False):
     if dim is None:
         dim = ()
-    return any(x, dim, keepdim)
+    return any(input, dim, keepdim)
 
 
 # 9 arctan2
@@ -948,6 +948,10 @@ def deprecated_tensor_remainder(input, divisor):
     return remainder(input, divisor)
 
 
+def deprecated_tensor_mod(input, other):
+    return _tensor_mod(input, other)
+
+
 # 86 repeat
 def tensor_repeat(input, *repeats):
     raise RuntimeError("'repeat' is not supported on this device.")
@@ -1514,6 +1518,16 @@ def tensor_empty(*size, dtype=None, device=None):
         "This is a function for empty not should be called. Please check the implementation.")
 
 
+def tensor_empty_like(input, *, dtype=None, device=None):
+    raise NotImplementedError(
+        "This is a function for empty_like should not be called. Please check the implementation.")
+
+
+def tensor_new_empty(input, size, *, dtype=None, device=None):
+    raise NotImplementedError(
+        "This is a function for new_empty should not be called. Please check the implementation.")
+
+
 def deprecated_tensor_logaddexp(input, other):
     return F.logaddexp(input, other)
 
@@ -1744,6 +1758,10 @@ def deprecated_tensor_diag(input):
     return F.diag(input)
 
 
+def deprecated_einsum(equation, operands):
+    raise NotImplementedError('einsum only supports Ascend.')
+
+
 # 916 index_add
 @constexpr
 def _check_index_add_alpha(alpha):
@@ -1808,6 +1826,8 @@ def tensor_inplace_sub(input, other, *, alpha=1):
         return sub(input, other)
     return sub_ext(input, other, alpha=alpha)
 
+def tensor_new_full(input, size, fill_value, *, dtype=None):
+    raise NotImplementedError("new_full method support Ascend only")
 
 def tensor_div_empty_(input, other, rounding_mode=None):
     raise ValueError("should not come here for div_ method.")
@@ -1840,6 +1860,10 @@ def all_gather_matmul(
     raise NotImplementedError('all_gather_matmul only supports Ascend.')
 
 
+def conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    raise NotImplementedError('conv3d only supports Ascend.')
+
+
 def tensor_floor_divide_(input, other):
     return _tensor_floordiv(input, other)
 
@@ -1888,6 +1912,14 @@ def deprecated_pixel_shuffle(input, upscale_factor):
     return F.pixel_shuffle(input, upscale_factor)
 
 
+def tensor_quant_matmul(x1, x2, scale, *, offset=None, pertoken_scale=None, bias=None, output_dtype=None,
+                        x1_dtype=None, x2_dtype=None, pertoken_scale_dtype=None, scale_dtype=None, group_sizes=None):
+    r"""
+    For details, please refer to :func:`mindspore.ops.auto_generate.quant_matmul`.
+    """
+    raise NotImplementedError('quant_matmul only supports Ascend.')
+
+
 def tensor_gmm(x, weight, *, bias=None, group_list=None, group_type=0, group_list_type=0):
     raise NotImplementedError("gmm has not been implemented by python.")
 

mindspore/ops_generate/api/cpp_create_prim_instance_helper_generator.py

@@ -53,7 +53,7 @@ class CppCreatePrimInstanceHelperGenerator(BaseGenerator):
         """
         py_arg_default = self.generate_op_arg_default_value(op_protos)
         py_labels = self.generate_op_labels(op_protos)
-        res_str = (template.PY_LICENSE_STR + py_arg_default + py_labels)
+        res_str = template.PY_LICENSE_STR + py_arg_default + py_labels
 
         save_path = os.path.join(work_path, K.PY_AUTO_GEN_PATH)
         file_name = "cpp_create_prim_instance_helper.py"

mindspore/ops_generate/api/functional_map_cpp_generator.py

@@ -159,14 +159,6 @@ class FunctionalMapCppGenerator(BaseGenerator):
         sig_str += '}\n},'
         return sig_str
 
-    def _is_input_arg(self, arg_name, op_name):
-        res = False
-        if op_name in K.INPUT_NAME_MAP and arg_name == K.INPUT_NAME_MAP[op_name]:
-            res = True
-        elif op_name not in K.INPUT_NAME_MAP and arg_name in K.INPUT_ARGS_NAME:
-            res = True
-        return res
-
     def _generate_single_signature_str(self, func_api_name, tensor_proto, is_tensor_method) -> str:
         """
         Generates a single function signature string for the given operation prototype.
@@ -186,7 +178,7 @@ class FunctionalMapCppGenerator(BaseGenerator):
         arg_valid_types = []
         for _, arg in enumerate(op_proto.op_args):
             arg_name = arg.arg_name
-            if is_tensor_method and self._is_input_arg(arg_name, op_name):
+            if is_tensor_method and _is_input_arg(arg_name, op_name):
                 continue
 
             arg_valid_types = self._handle_arg_valid_types(arg, arg_name, arg_valid_types, func_api_name)
@@ -502,3 +494,12 @@ class FunctionalMapCppGenerator(BaseGenerator):
                 self._get_and_append_single_op_varargs_list(func_protos,
                                                             mint_varargs_list)
         return mint_varargs_list
+
+
+def _is_input_arg(arg_name, op_name):
+    res = False
+    if op_name in K.INPUT_NAME_MAP and arg_name == K.INPUT_NAME_MAP[op_name]:
+        res = True
+    elif op_name not in K.INPUT_NAME_MAP and arg_name in K.INPUT_ARGS_NAME:
+        res = True
+    return res

mindspore/ops_generate/api/functions_cc_generator.py

@@ -37,6 +37,13 @@ class FunctionsHeaderGenerator(BaseGenerator):
         """
         self.FUNCTIONS_H_TEMPLATE = template.FUNCTIONS_H_TEMPLATE
         self.function_interface_template = Template("${return_type} BACKEND_EXPORT ${op_name}(${input_args});")
+        self.function_interface_template_comm = Template(
+            "${return_type} BACKEND_EXPORT ${op_name}_inner(${input_args}," \
+            "CommHandlePtr comm_handle, const std::string& target);"
+        )
+        self.function_interface_template_comm_return_handle = Template(
+            "${return_type_with_handle} BACKEND_EXPORT ${op_name}(${input_args});"
+        )
 
     def generate(self, work_path, op_protos):
         """
@@ -48,14 +55,30 @@ class FunctionsHeaderGenerator(BaseGenerator):
         """
         functions_list = []
         for op_proto in op_protos:
-            if op_proto.op_dispatch is None or op_proto.op_dispatch.is_comm_op:
+            if op_proto.op_dispatch is None:
                 continue
             input_args_with_type_str = self._get_input_args(op_proto)
             return_type_str = _get_return_type_str(op_proto)
-            functions = self.function_interface_template.replace(op_name=op_proto.op_name,
-                                                                 input_args=input_args_with_type_str,
-                                                                 return_type=return_type_str)
+            function_template = (
+                self.function_interface_template
+                if not op_proto.op_dispatch.is_comm_op
+                else self.function_interface_template_comm
+            )
+            functions = function_template.replace(op_name=op_proto.op_name,
+                                                  input_args=input_args_with_type_str,
+                                                  return_type=return_type_str)
             functions_list.append(functions)
+            if op_proto.op_dispatch.is_comm_op:
+                return_type_with_handle = _get_return_type_with_handle_str(return_type_str)
+                functions_with_handle = (
+                    self.function_interface_template_comm_return_handle.replace(
+                        op_name=op_proto.op_name,
+                        input_args=input_args_with_type_str,
+                        return_type=return_type_str,
+                        return_type_with_handle=return_type_with_handle,
+                    )
+                )
+                functions_list.append(functions_with_handle)
         pyboost_func_h_str = self.FUNCTIONS_H_TEMPLATE.replace(op_call_with_grad=functions_list)
         save_path = os.path.join(work_path, K.MS_PYBOOST_FUNCTIONS_AUTO_GEN_PATH)
         file_name = "functions.h"
@@ -73,7 +96,7 @@ class FunctionsHeaderGenerator(BaseGenerator):
         """
         args_list = []
         for op_arg in op_proto.op_args:
-            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg))
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg), op_proto.op_view)
             args_list.append("const " + input_dtype + " &" + op_arg.arg_name)
         return args_list
 
@@ -89,12 +112,19 @@ class FunctionsGenerator(BaseGenerator):
         """
         self.FUNCTIONS_CC_TEMPLATE = template.FUNCTIONS_CC_TEMPLATE
         self.FUNCTION_BODY_TEMPLATE = template.FUNCTION_BODY_TEMPLATE
+        self.FUNCTION_COMM_BODY_TEMPLATE = template.FUNCTION_COMM_BODY_TEMPLATE
         self.pyboost_func_include_header_template = Template(
             f'#include "{K.MS_PYBOOST_BASE_PATH}/auto_generate/${{operator_name}}.h"\n'
         )
         self.clone_inplace_input_template = Template(
             'GetCloneFunc()(op, prim::kPrim${class_name}, device_target, {${grad_args}});'
         )
+        self.create_aclnn_op_template = Template(
+            'auto op = CREATE_PYBOOST_OP(${class_name}, device_target);'
+        )
+        self.create_internal_op_template = Template(
+            'auto op = CREATE_PYBOOST_SELECTED_OP(${class_name}, device_target);'
+        )
 
     def generate(self, work_path, op_protos):
         """
@@ -108,7 +138,7 @@ class FunctionsGenerator(BaseGenerator):
         op_call_with_grad_list = []
         ops_inc_head_set = set()
         for op_proto in op_protos:
-            if op_proto.op_dispatch is None or op_proto.op_dispatch.is_comm_op:
+            if op_proto.op_dispatch is None:
                 continue
             func_include_headers_list.append(
                 self.pyboost_func_include_header_template.replace(operator_name=op_proto.op_name))
@@ -137,8 +167,22 @@ class FunctionsGenerator(BaseGenerator):
         inplace_clone_args = self._get_clone_input_args(op_proto, False, False)
         clone_func_str = self._get_clone_inplace_str(op_proto.op_inplace, op_proto.op_class.name, inplace_clone_args)
         return_type_str = _get_return_type_str(op_proto)
+        if op_proto.op_dispatch.is_comm_op:
+            return_type_with_handle = _get_return_type_with_handle_str(return_type_str)
+            comm_body = self.FUNCTION_COMM_BODY_TEMPLATE.replace(op_name=op_proto.op_name,
+                                                                 class_name=op_proto.op_class.name,
+                                                                 input_args=input_args,
+                                                                 clone_func=clone_func_str,
+                                                                 input_args_with_type=input_args_with_type,
+                                                                 return_type=return_type_str,
+                                                                 return_type_with_handle=return_type_with_handle)
+            return comm_body
+        create_op_str = self.create_aclnn_op_template.replace(class_name=op_proto.op_class.name)
+        if getattr(op_proto.op_dispatch, 'internal_op_ascend') != 'None':
+            create_op_str = self.create_internal_op_template.replace(class_name=op_proto.op_class.name)
         return self.FUNCTION_BODY_TEMPLATE.replace(op_name=op_proto.op_name,
                                                    class_name=op_proto.op_class.name,
+                                                   create_op=create_op_str,
                                                    input_args=input_args,
                                                    clone_func=clone_func_str,
                                                    input_args_with_type=input_args_with_type,
@@ -157,7 +201,7 @@ class FunctionsGenerator(BaseGenerator):
         """
         args_list = []
         for op_arg in op_proto.op_args:
-            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg))
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg), op_proto.op_view)
             if has_type:
                 args_list.append("const " + input_dtype + " &" + op_arg.arg_name)
             else:
@@ -194,7 +238,7 @@ class FunctionsGenerator(BaseGenerator):
         """
         args_list = []
         for op_arg in op_proto.op_args:
-            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg))
+            input_dtype = get_input_dtype(op_arg.arg_dtype, is_optional_param(op_arg), op_proto.op_view)
             if has_type:
                 args_list.append(f"const {input_dtype} &{op_arg.arg_name}")
             else:
@@ -205,6 +249,10 @@ class FunctionsGenerator(BaseGenerator):
         return args_list
 
 
+def _get_return_type_with_handle_str(return_type_str):
+    return f"std::tuple<{return_type_str}, CommHandlePtr>"
+
+
 def _get_return_type_str(op_proto):
     """
     Get the return type string for the function.
@@ -217,8 +265,8 @@ def _get_return_type_str(op_proto):
     """
     returns_type = []
     type_convert_to_base = {
-        'std::vector<mindspore::tensor::TensorPtr>': 'std::vector<mindspore::tensor::BaseTensorPtr>',
-        'mindspore::tensor::TensorPtr': 'mindspore::tensor::BaseTensorPtr'
+        'std::vector<mindspore::tensor::TensorPtr>': 'std::vector<mindspore::tensor::TensorPtr>',
+        'mindspore::tensor::TensorPtr': 'mindspore::tensor::TensorPtr'
     }
     for return_obj in op_proto.op_returns:
         temp_return = get_return_type(return_obj.arg_dtype)

mindspore/ops_generate/api/tensor_func_reg_cpp_generator.py

@@ -81,7 +81,7 @@ class TensorFuncRegCppGenerator(BaseGenerator):
             'MS_LOG(INFO) << "Call Tensor${class_name}";\n'
             'auto res = mindspore::pynative::'
             '${pyboost_function}(mindspore::prim::kPrim${class_name}, parse_args.src_types_, ${convert_args});\n'
-            'trace::Capture(parse_args.arg_list_, "${class_name}", &res);\n'
+            'trace::Capture(parse_args.arg_list_, mindspore::prim::kPrim${class_name}, &res);\n'
             'return res;\n'
         )
         self.callback_python_template = Template(

mindspore/ops_generate/common/base_generator.py

@@ -1,3 +1,17 @@
+# 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.
+# ============================================================================
 """
 Base generator interface for other pyboost file generator classes
 """