mindspore 2.7.0__cp311-cp311-win_amd64.whl → 2.7.1__cp311-cp311-win_amd64.whl

mindspore/ops/operations/array_ops.py

@@ -1157,6 +1157,7 @@ class FillV2(PrimitiveWithCheck):
             init_func = Zero()
             init_func.__enable_zero_dim__ = True
             out = Tensor(shape=dims, dtype=x.dtype, init=init_func)
+            out.init_data()
             return out
         return Tensor(np.full(dims, x.asnumpy()))
 

mindspore/ops/operations/comm_ops.py

@@ -1940,7 +1940,7 @@ class BatchISendIRecv(PrimitiveWithInfer):
 
 
 class AlltoAllV(PrimitiveWithInfer):
-    """
+    r"""
     AllToAllV which support uneven scatter and gather compared with AllToAll.
 
     Note:
@@ -2001,7 +2001,7 @@ class AlltoAllV(PrimitiveWithInfer):
         ...    send_tensor = Tensor([0, 1, 2.])
         ...    send_numel_list = [1, 2]
         ...    recv_numel_list = [1, 2]
-        >>> elif rank == 1:
+        ... elif rank == 1:
         ...    send_tensor = Tensor([3, 4, 5.])
         ...    send_numel_list = [2, 1]
         ...    recv_numel_list = [2, 1]
@@ -2013,6 +2013,10 @@ class AlltoAllV(PrimitiveWithInfer):
         rank 1:
         [1. 2. 5]
 
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - AlltoAllV
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#alltoallv>`_
+
     """
 
     @prim_attr_register
@@ -2024,6 +2028,94 @@ class AlltoAllV(PrimitiveWithInfer):
         self.add_prim_attr('block_size', self.block_size)
 
 
+class AlltoAllVC(PrimitiveWithInfer):
+    r"""
+    AllToAllVC passes in the sending and receiving parameters of all ranks through the input parameter
+    `send_count_matrix`. Compared to AllToAllV, AllToAllVC does not require the aggregation of all rank
+    sending and receiving parameters, thus offering superior performance.
+
+    Note:
+        Only one-dimensional input is supported; the input data must be flattened into a one-dimensional
+        array before using this interface.
+
+    Args:
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``, which
+            means ``"hccl_world_group"`` in Ascend.
+        block_size (int, optional): The basic units for splitting and gathering numel by `send_count_matrix`.
+            Default: ``1``.
+        transpose (bool, optional): Indicates whether the `send_count_matrix` needs to undergo a transpose
+            operation, this parameter is used in reverse calculation scenarios. Default: ``False``.
+
+    Inputs:
+        - **input_x** (Tensor) - flatten tensor to scatter. The shape of tensor is :math:`(x_1)`.
+        - **send_count_matrix** (Union[list[int], Tensor]) - The sending and receiving parameters of
+          all ranks, :math:`\text{send_count_matrix}[i*\text{rank_size}+j]` represents the amount of data sent by
+          rank i to rank j, and the basic unit is the number of bytes of Tensor's dtype. Among them, `rank_size`
+          indicates the size of the communication group.
+
+    Outputs:
+        Tensor. Flattened and concatenated tensor gather from remote ranks.
+        If gather result is empty, it will return a Tensor with shape `()`, and value has no actual meaning.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend/GPU/CPU devices, it is recommended to use the msrun startup method
+            without any third-party or configuration file dependencies.
+
+            Please see the `msrun start up
+            <https://www.mindspore.cn/tutorials/en/master/parallel/msrun_launcher.html>`_
+            for more details.
+
+            This example should be run with 2 devices.
+
+        >>> from mindspore.ops import AlltoAllVC
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init, get_rank
+        >>> from mindspore import Tensor
+        >>>
+        >>> init()
+        >>> rank = get_rank()
+        >>> class Net(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net, self).__init__()
+        ...         self.all_to_all_v_c = AlltoAllVC()
+        ...
+        ...     def construct(self, x, send_count_matrix):
+        ...         return self.all_to_all_v_c(x, send_count_matrix)
+        >>> send_count_matrix = Tensor([[0, 3], [3, 0]])
+        >>> send_tensor = Tensor([0, 1, 2.]) * rank
+        >>> net = Net()
+        >>> output = net(send_tensor, send_count_matrix)
+        >>> print(output)
+        rank 0:
+        [0. 1. 2]
+        rank 1:
+        [0. 0. 0]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - AlltoAllVC
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#alltoallvc>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, group=GlobalComm.WORLD_COMM_GROUP, block_size=1, transpose=False):
+        self.group = GlobalComm.WORLD_COMM_GROUP if group is None else _get_group(group)
+        self.rank_size = get_group_size(self.group)
+        self.add_prim_attr('rank_size', self.rank_size)
+        self.add_prim_attr('group', self.group)
+        self.rank_id = get_rank(_get_group(self.group))
+        self.add_prim_attr('rank_id', self.rank_id)
+        validator.check_value_type("block_size", block_size, [int], self.name)
+        self.add_prim_attr('block_size', self.block_size)
+        self.add_prim_attr('transpose', self.transpose)
+
+
 class AllGatherV(PrimitiveWithInfer):
     """
     Gathers uneven tensors from the specified communication group and returns the tensor which is all gathered.

mindspore/ops/operations/custom_ops.py

@@ -18,6 +18,7 @@ from __future__ import absolute_import
 import json
 import os
 import re
+import sys
 import ast
 import hashlib
 import stat
@@ -26,6 +27,7 @@ import inspect
 import importlib
 import platform
 import subprocess
+import shutil
 import numpy as np
 import mindspore as ms
 from mindspore._c_expression import Oplib, typing
@@ -37,6 +39,7 @@ from mindspore.ops import DataType
 from mindspore import log as logger
 from mindspore import ops
 from mindspore.communication.management import get_rank, GlobalComm
+from mindspore import _checkparam as validator
 from ._ms_kernel import determine_variable_usage
 from ._custom_grad import autodiff_bprop
 from ._pyfunc_registry import add_pyfunc
@@ -1185,6 +1188,54 @@ class Custom(ops.PrimitiveWithInfer):
         return ops.primitive._run_op(self, self.name, args)
 
 
+class _MultiSoProxy:
+    """
+    A thin wrapper that transparently multiplexes attribute access between a
+    pure-Python fallback module and an optional compiled shared-object (SO)
+    module, honoring MindSpore’s current execution mode (GRAPH vs. PYNATIVE).
+    """
+
+    def __init__(self, func_module, so_module):
+        """
+        Args:
+            func_module (module or None): Python module to serve as the fallback implementation source.
+                May be ``None`` if no Python fallback is available.
+            so_module (module): Compiled shared-object module that provides
+                optimized kernels accessible only in ``PYNATIVE_MODE``.
+        """
+        self.func_module = func_module
+        self.so_module = so_module
+
+    def __getattr__(self, name: str):
+        """
+        Intercepts every attribute lookup and resolves it against the wrapped
+        modules according to the documented precedence rules.
+
+        Args:
+            name (str): Name of the custom operation being requested.
+
+        Returns:
+            Any: The requested callable or attribute from either ``func_module`` or ``so_module``.
+
+        Raises:
+            AttributeError: If the attribute is not found in any applicable module or
+            is incompatible with the current execution mode.
+        """
+        if self.func_module is not None and hasattr(self.func_module, name):
+            return getattr(self.func_module, name)
+        if context.get_context("mode") == ms.PYNATIVE_MODE:
+            if hasattr(self.so_module, name):
+                return getattr(self.so_module, name)
+            raise AttributeError(
+                f"Custom op '{name}' is neither in func_module nor in so_module."
+            )
+
+        raise AttributeError(
+            f"Custom op '{name}' does not support GRAPH mode. "
+            f"Please register it for GRAPH mode or switch to PYNATIVE mode."
+        )
+
+
 class CustomOpBuilder:
     r"""
     CustomOpBuilder is used to initialize and configure custom operators for MindSpore.
@@ -1200,10 +1251,11 @@ class CustomOpBuilder:
 
     Args:
         name (str): The unique name of the custom operator module, used to identify the operator.
-        sources (Union[str, list[str]]): The source file(s) of the custom operator. It can be a single file path or
-                                    a list of file paths.
+        sources (Union[list[str], tuple[str], str]): The source file(s) of the custom operator. It can be a single
+            file path or a list of file paths.
         backend (str, optional): The target backend for the operator, such as "CPU" or "Ascend". Default: ``None``.
-        include_paths (list[str], optional): Additionally included paths needed during compilation. Default: ``None``.
+        include_paths (Union[list[str], tuple[str], str], optional): Additionally included paths needed during
+            compilation. Default: ``None``.
         cflags (str, optional): Extra C++ compiler flags to be used during compilation. Default: ``None``.
         ldflags (str, optional): Extra linker flags to be used during linking. Default: ``None``.
         kwargs (dict, optional): Additional keyword arguments for future extensions or specific custom requirements.
@@ -1217,6 +1269,17 @@ class CustomOpBuilder:
             - enable_atb (bool, optional): Whether to call ATB (Ascend Transformer Boost) operator. If set to ``True``,
               the `backend` must be ``Ascend`` or left empty. Default: ``False``.
 
+            - enable_asdsip (bool, optional): Whether to call ASDSIP (Ascend SiP Boost) operator. If set to ``True``,
+              the `backend` must be ``Ascend`` or left empty. Default: ``False``.
+
+            - op_def (Union[list[str], tuple[str], str], optional): Path(s) to the operator definition
+              file(s) (YAML format). When using custom operators in graph mode, this parameter is mandatory.
+              It can be a single file path string or a list of file path strings. Default: ``None``.
+
+            - op_doc (Union[list[str], tuple[str], str], optional): Path(s) to the operator documentation
+              file(s) (YAML format). This parameter is optional and used to provide additional documentation
+              for the operator. It can be a single file path string or a list of file path strings. Default: ``None``.
+
     .. note::
         - If the `backend` argument is provided, additional default flags will be automatically added to
           the compilation and linking steps to support the operator's target backend. The default options
@@ -1239,21 +1302,20 @@ class CustomOpBuilder:
     _loaded_ops = {}
 
     def __init__(self, name, sources, backend=None, include_paths=None, cflags=None, ldflags=None, **kwargs):
-        self.name = name
-        self.source = sources
-        self.backend = backend
-        self.include_paths = include_paths
-        self.cflags = cflags
-        self.ldflags = ldflags
-        self.build_dir = kwargs.get("build_dir")
-        self.enable_atb = kwargs.get("enable_atb", False)
-        self.debug_mode = kwargs.get("debug_mode", False)
+        self._check_and_get_args(name, sources, backend, include_paths, cflags, ldflags, **kwargs)
+
         self._ms_path = os.path.dirname(os.path.abspath(ms.__file__))
+        self.auto_generate = self.name + "_auto_generate"
         if self.enable_atb:
             if backend is not None and backend != "Ascend":
                 raise ValueError("For 'CustomOpBuilder', when 'enable_atb' is set to True, the 'backend' must be "
                                  f"'Ascend' (or left implicit), but got '{backend}'")
             self.backend = "Ascend"
+        if self.enable_asdsip:
+            if backend is not None and backend != "Ascend":
+                raise ValueError("For 'CustomOpBuilder', when 'enable_asdsip' is set to True, the 'backend' must be "
+                                 f"'Ascend' (or left implicit), but got '{backend}'")
+            self.backend = "Ascend"
         if self.backend == "Ascend":
             ascend_opp_path = os.getenv("ASCEND_OPP_PATH")
             if not ascend_opp_path:
@@ -1265,6 +1327,115 @@ class CustomOpBuilder:
                 if not self.atb_home_path:
                     raise ValueError("Environment variable 'ATB_HOME_PATH' must be set when 'enable_atb' is True.")
 
+    def _check_and_get_args(self, name, sources, backend=None, include_paths=None,
+                            cflags=None, ldflags=None, **kwargs):
+        """
+        Validate and normalize all arguments to meet custom-op build requirements.
+        """
+
+        def _check_str_or_list_str(key, val):
+            if val is None:
+                return val
+            if isinstance(val, str):
+                val = [val]
+            val = validator.check_value_type(key, val, [list, tuple])
+            val = list(val)
+            validator.check_element_type_of_iterable(key, val, [str])
+            return val
+
+        self.name = validator.check_value_type("name", name, [str])
+        self.source = _check_str_or_list_str("sources", sources)
+        self.backend = validator.check_value_type("backend", backend, [str, type(None)])
+        if self.backend is not None and self.backend not in {"CPU", "Ascend"}:
+            raise ValueError(
+                f"For 'backend', only 'CPU' or 'Ascend' are allowed, but got '{self.backend}'.")
+
+        self.include_paths = _check_str_or_list_str("include_paths", include_paths)
+
+        self.cflags = validator.check_value_type("cflags", cflags, [str, type(None)])
+        self.ldflags = validator.check_value_type("ldflags", ldflags, [str, type(None)])
+
+        self.build_dir = validator.check_value_type("build_dir",
+                                                    kwargs.get("build_dir"),
+                                                    [str, type(None)])
+
+        self.debug_mode = validator.check_bool(kwargs.get("debug_mode", False), "debug_mode")
+        self.enable_asdsip = validator.check_bool(kwargs.get("enable_asdsip", False), "enable_asdsip")
+        self.yaml = _check_str_or_list_str("op_def", kwargs.get("op_def"))
+        self.doc = _check_str_or_list_str("op_doc", kwargs.get("op_doc"))
+
+        self.enable_atb = validator.check_bool(kwargs.get("enable_atb", False))
+
+    def _generate_custom_op_def(self, module: str, input_path: str, doc_path: str, output_path: str) -> None:
+        """Call gen_custom_ops.py to generate custom operator definition"""
+        file_path = os.path.join(self._ms_path, "ops_generate/gen_custom_ops.py")
+        cmd = [
+            sys.executable,
+            file_path,
+            "-i", input_path,
+            "-o", output_path,
+            "-m", module,
+            "-d", doc_path
+        ]
+
+        try:
+            subprocess.run(
+                cmd,
+                check=True,
+                text=True,
+                capture_output=True
+            )
+        except subprocess.CalledProcessError as exc:
+            raise RuntimeError(
+                f"gen_custom_op.py failed with exit code {exc.returncode}.\n"
+                f"stdout: {exc.stdout}\n"
+                f"stderr: {exc.stderr}"
+            ) from None
+
+    def _get_op_def(self):
+        """
+        Generate C++ operator-definition source files from one or more YAML specification files.
+        """
+        if self.yaml is None:
+            return []
+
+        if self.doc is None:
+            logger.info("Missing required 'doc': no YAML document was provided.")
+
+        build_path = self._get_build_directory()
+        yaml_path = os.path.join(build_path, "yaml")
+        op_def_path = os.path.join(build_path, self.auto_generate)
+        if os.path.exists(op_def_path):
+            shutil.rmtree(op_def_path)
+        os.makedirs(op_def_path, exist_ok=True)
+
+        def copy_files(yaml_files, dest_path):
+            if os.path.exists(dest_path):
+                shutil.rmtree(dest_path)
+            os.makedirs(dest_path, exist_ok=True)
+            for file_path in yaml_files:
+                if not os.path.isfile(file_path):
+                    raise FileNotFoundError(f"File not found: {file_path}")
+
+                filename = os.path.basename(file_path)
+                file_ext = os.path.splitext(filename)[1].lower()
+                if file_ext not in ('.yaml', '.yml'):
+                    raise ValueError(f"Invalid file extension: {file_ext} for {filename}")
+
+                _dest_path = os.path.join(dest_path, filename)
+                shutil.copy2(file_path, _dest_path)
+
+        yaml_files = [self.yaml] if isinstance(self.yaml, str) else self.yaml
+        copy_files(yaml_files, yaml_path)
+        doc_path = ""
+        if self.doc is not None:
+            doc_path = os.path.join(build_path, "doc")
+            doc_files = [self.doc] if isinstance(self.doc, str) else self.doc
+            copy_files(doc_files, doc_path)
+
+        self._generate_custom_op_def(self.name, yaml_path, doc_path, op_def_path)
+        return [os.path.join(op_def_path, "gen_custom_ops_def.cc")]
+
     def get_sources(self):
         """
         Get the source files for the custom operator.
@@ -1272,7 +1443,8 @@ class CustomOpBuilder:
         Returns:
             str or list[str], The source file(s) for the operator.
         """
-        return self.source
+        self.source = [self.source] if isinstance(self.source, str) else self.source
+        return self.source + self._get_op_def()
 
     def get_include_paths(self):
         """
@@ -1299,6 +1471,7 @@ class CustomOpBuilder:
         """include paths for inner module interface."""
         ms_inner_path = os.path.join(self._ms_path, "include", "mindspore")
         include_list = []
+        include_list.append(os.path.join(ms_inner_path, "include"))
         include_list.append(os.path.join(ms_inner_path, "core", "include"))
         include_list.append(os.path.join(ms_inner_path, "core", "mindrt", "include"))
         include_list.append(os.path.join(ms_inner_path, "core", "mindrt"))
@@ -1320,10 +1493,14 @@ class CustomOpBuilder:
         flags += ['-std=c++17', '-fstack-protector-all', '-fPIC', '-pie']
         if self.debug_mode:
             flags.append('-g')
+        else:
+            flags.append('-O2')
         if self.backend == "Ascend":
             flags.append('-DCUSTOM_ASCEND_OP')
             if self.enable_atb:
                 flags.append('-DCUSTOM_ENABLE_ATB')
+            if self.enable_asdsip:
+                flags.append('-DCUSTOM_ENABLE_ASDSIP')
         if self.cflags is not None:
             flags.append(self.cflags)
         return flags
@@ -1344,18 +1521,23 @@ class CustomOpBuilder:
             '-lmindspore_core',
             '-lmindspore_ms_backend',
             '-lmindspore_pynative',
-            '-lmindspore_extension'
+            '-lmindspore_pyboost'
         ]
         if self.backend == "Ascend":
-            flags.append(f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin'))}")
             flags.append(f"-L{os.path.abspath(os.path.join(self.ascend_cann_path, 'lib64'))}")
             flags.append('-lascendcl')
+            plugin_path = os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin'))
+            flags.append(f"-L{plugin_path}")
+            flags.append(f"-L{os.path.join(plugin_path, 'ascend')}")
             flags.append('-l:libmindspore_ascend.so.2')
+            flags.append('-lmindspore_extension_ascend_aclnn')
             if self.enable_atb:
-                flags.append(f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin', 'ascend'))}")
                 flags.append('-lmindspore_extension_ascend_atb')
                 flags.append(f"-L{os.path.abspath(os.path.join(self.atb_home_path, 'lib'))}")
                 flags.append('-latb')
+            if self.enable_asdsip:
+                flags.append(f"-L{os.path.abspath(os.path.join(self._ms_path, 'lib', 'plugin', 'ascend'))}")
+                flags.append('-lmindspore_extension_ascend_asdsip')
         if self.ldflags is not None:
             flags.append(self.ldflags)
         return flags
@@ -1386,15 +1568,42 @@ class CustomOpBuilder:
         """
         if self.name in CustomOpBuilder._loaded_ops:
             return CustomOpBuilder._loaded_ops[self.name]
+
         module_path = self.build()
-        mod = self._import_module(module_path)
+        so_module = CustomOpBuilder._import_module(module_path)
+        func_module = None
+        if self.yaml is not None:
+            module_path = os.path.join(self.build_dir, self.auto_generate, "gen_ops_def.py")
+            sys.path.append(os.path.join(self.build_dir, self.auto_generate))
+            sys.path.append(os.path.join(self.build_dir))
+            func_module = self._import_module(module_path, True)
+        mod = _MultiSoProxy(func_module, so_module)
         CustomOpBuilder._loaded_ops[self.name] = mod
         return mod
 
-    def _import_module(self, module_path):
+    @staticmethod
+    def _import_module(module_path, is_yaml_build=False):
         """Import module from library."""
-        spec = importlib.util.spec_from_file_location(self.name, module_path)
+        module_path = os.path.abspath(module_path)
+        module_dir = os.path.dirname(module_path)
+        module_name = os.path.splitext(os.path.basename(module_path))[0]
+
+        if is_yaml_build:
+            package_name = os.path.basename(module_dir)
+            if module_dir not in sys.path:
+                sys.path.append(module_dir)
+
+            if package_name not in sys.modules:
+                pkg_spec = importlib.machinery.ModuleSpec(package_name, None, is_package=True)
+                pkg = importlib.util.module_from_spec(pkg_spec)
+                pkg.__path__ = [module_dir]
+                sys.modules[package_name] = pkg
+
+            module_name = f"{package_name}.{module_name}"
+
+        spec = importlib.util.spec_from_file_location(module_name, module_path)
         module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
         spec.loader.exec_module(module)
         return module
 

mindspore/ops/operations/debug_ops.py

@@ -13,14 +13,11 @@
 # limitations under the License.
 # ============================================================================
 """debug_ops"""
-import stat
-from pathlib import Path
-
-import numpy as np
+import inspect
 from mindspore import log as logger
 from mindspore._c_expression import security, HookType
 from mindspore._c_expression import TensorPy as Tensor_
-from mindspore._c_expression import _tensordump_process_file
+from mindspore._c_expression import _tensordump_exec
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
@@ -314,26 +311,12 @@ class TensorDump(Primitive):
         self.add_prim_attr("side_effect_io", True)
         self.add_prim_attr("channel_name", "ms_tensor_dump")
 
-    def _save_file(self, file, data):
-        file = Path(file)
-        if file.exists():
-            file.chmod(stat.S_IWUSR)
-        np.save(file, data)
-        file.chmod(stat.S_IRUSR)
-
     def __call__(self, file, input_x):
         validator.check_value_type('file', file, [str], self.__class__.__name__)
         if not file:
             raise ValueError("For 'TensorDump', the input argument[file] cannot be an empty string.")
         validator.check_value_type('input_x', input_x, [Tensor], self.__class__.__name__)
-
-        dtype = input_x.dtype
-        file = _tensordump_process_file(file, str(dtype))
-        if not file:
-            return
-        if dtype == mstype.bfloat16:
-            input_x = P.Cast()(input_x, mstype.float32)
-        self._save_file(file, input_x.asnumpy())
+        _tensordump_exec(file, input_x)
 
 
 class HistogramSummary(Primitive):
@@ -529,14 +512,15 @@ class Morph(PrimitiveWithInfer):
 
     .. note::
         - This primitive is only supported in GRAPH_MODE.
-        - `fn` must satisfy the syntax constraints of the graph mode.
-        - Users do not need to implement a custom backward function.
+        - A user-defined bprop (by argument: `bprop_fn`) is allowed for `Morph`.
+        - `fn` and `bprop_fn` must satisfy the syntax constraints of the graph mode.
         - `vararg`, `kwarg`, `kwonlyargs` and free variables are not supported in user-defined function.
 
     Args:
-        fn (Function): Mindspore's function, user-defined function.
-        infer_shape (Function): Mindspore's function, user-defined infer_shape function.
-        infer_dtype (Function): Mindspore's function, user-defined infer_dtype function.
+        fn (Function): MindSpore's function, user-defined function.
+        infer_shape (Function): MindSpore's function, user-defined infer_shape function.
+        infer_dtype (Function): MindSpore's function, user-defined infer_dtype function.
+        bprop_fn (Function, optional): MindSpore's function, user-defined bprop function, default: ``None``.
 
     Inputs:
         The inputs of user-defined `fn`.
@@ -590,21 +574,35 @@ class Morph(PrimitiveWithInfer):
         >>> weight0_grad = bwd_out[1][0].asnumpy()
         >>> weight1_grad = bwd_out[1][1].asnumpy()
         >>> print("x_grad", x_grad)
-        >>> print("weight0_grad", weight0_grad)
-        >>> print("weight1_grad", weight1_grad)
         x_grad [ 400. 1000. 1800.]
+        >>> print("weight0_grad", weight0_grad)
         weight0_grad [2800. 4000. 5400.]
+        >>> print("weight1_grad", weight1_grad)
         weight1_grad [ 700. 1600. 2700.]
     """
     @prim_attr_register
-    def __init__(self, fn, infer_shape, infer_dtype):
+    def __init__(self, fn, infer_shape, infer_dtype, bprop_fn=None):
         self.add_prim_attr('side_effect_backprop', True)
         self.add_prim_attr('side_effect_mem', True)
         self.add_prim_attr('side_effect_io', True)
-        self.add_prim_attr('__metamorphosis__', fn)
         self._infer_shape = infer_shape
         self._infer_dtype = infer_dtype
 
+        self.add_prim_attr('__metamorphosis__', True)
+        self.__morph_fn__ = fn
+        self.__morph_bprop_fn__ = None
+        if bprop_fn:
+            self._check_fn_supported(fn)
+            self.__morph_bprop_fn__ = bprop_fn
+
+    def _check_fn_supported(self, fn):
+        fn_sig = inspect.signature(fn)
+        for param in fn_sig.parameters.values():
+            if not (param.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD and param.default is inspect.Parameter.empty):
+                raise ValueError(f"When use `bprop` in Morph, Morph `fn` only support positional or keyword parameters "
+                                 f"with default value is empty, but got param '{param.name}' "
+                                 f"of kind '{param.kind.name}' with default value '{param.default}'.")
+
     def infer_shape(self, *args):
         return self._infer_shape(*args)