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

mindspore/ops/operations/comm_ops.py

@@ -54,7 +54,7 @@ class ReduceOp:
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with multiple devices.
@@ -141,7 +141,7 @@ class AllReduce(Primitive):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -178,14 +178,15 @@ class AllReduce(Primitive):
     @prim_attr_register
     def __init__(self, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
         """Initialize AllReduce."""
+        self.group = _get_group(group)
         if not isinstance(op, type(ReduceOp.SUM)):
             raise TypeError(f"For '{self.name}', the 'op' must be str, but got {type(op).__name__}.")
-        if not isinstance(_get_group(group), str):
+        if not isinstance(self.group, str):
             raise TypeError(f"For '{self.name}', the 'group' must be str, "
-                            f"but got {type(_get_group(group)).__name__}.")
-        check_hcom_group_valid(group, prim_name=self.name)
+                            f"but got {type(self.group).__name__}.")
+        check_hcom_group_valid(self.group, prim_name=self.name)
         self.op = op
-        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('index', 0)
         self.add_prim_attr('no_eliminate', True)
@@ -230,7 +231,7 @@ class Reduce(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 4 devices.
@@ -314,7 +315,7 @@ class AllGather(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -354,12 +355,13 @@ class AllGather(PrimitiveWithInfer):
     @prim_attr_register
     def __init__(self, group=GlobalComm.WORLD_COMM_GROUP):
         """Initialize AllGather."""
-        validator.check_value_type('group', _get_group(group), (str,), self.name)
-        self.rank = get_rank(_get_group(group))
-        self.rank_size = get_group_size(_get_group(group))
+        self.group = _get_group(group)
+        validator.check_value_type('group', self.group, (str,), self.name)
+        self.rank = get_rank(self.group)
+        self.rank_size = get_group_size(self.group)
         validator.check('rank', self.rank, 'rank_size', self.rank_size, validator.LT, self.name)
         self.add_prim_attr('rank_size', self.rank_size)
-        self.add_prim_attr('group', _get_group(group))
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('mean_flag', False)
         self.add_prim_attr('no_eliminate', True)
@@ -375,25 +377,6 @@ class AllGather(PrimitiveWithInfer):
         return x_dtype
 
 
-class AShardIdentity(PrimitiveWithInfer):
-    """
-    Auto parallel virtual operator. Identity operator only for shard function.
-    Do nothing in terms of infer_shape, infer_dtype, and the tensor.
-
-    It is only for internal use of parallel modules and cannot be called by users.
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        pass
-
-    def infer_shape(self, x_shape):
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        return x_dtype
-
-
 class _MiniStepAllGather(PrimitiveWithInfer):
     """
     Auto parallel virtual operator. Do nothing in forward, do reducescatter in backward in mini-step. It is only for
@@ -555,7 +538,7 @@ class ReduceScatter(Primitive):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -597,11 +580,12 @@ class ReduceScatter(Primitive):
     def __init__(self, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
         """Initialize ReduceScatter."""
         validator.check_value_type('op', op, (type(ReduceOp.SUM),), self.name)
-        validator.check_value_type('group', _get_group(group), (str,), self.name)
+        self.group = _get_group(group)
+        validator.check_value_type('group', self.group, (str,), self.name)
         self.op = op
-        self.rank_size = get_group_size(_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', _get_group(group))
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('no_eliminate', True)
 
@@ -692,7 +676,7 @@ class Broadcast(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -922,7 +906,7 @@ class AlltoAll(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 8 devices.
@@ -1041,7 +1025,7 @@ class NeighborExchangeV2(Primitive):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1158,7 +1142,7 @@ class CollectiveScatter(Primitive):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1243,7 +1227,7 @@ class CollectiveGather(Primitive):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 4 devices.
@@ -1308,8 +1292,6 @@ class Barrier(PrimitiveWithInfer):
     Raises:
         TypeError: If `group` is not a str.
         RuntimeError: If backend is invalid, or distributed initialization fails.
-        ValueError: If the local rank id of the calling process in the group
-                    is larger than the group's rank size.
 
     Supported Platforms:
         ``Ascend``
@@ -1321,7 +1303,7 @@ class Barrier(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1395,7 +1377,7 @@ class Send(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1431,7 +1413,7 @@ class Send(PrimitiveWithInfer):
     def __init__(self, sr_tag, dest_rank, group=GlobalComm.WORLD_COMM_GROUP, group_back=GlobalComm.WORLD_COMM_GROUP):
         self.rank = dest_rank
         self.sr_tag = sr_tag
-        self.group = group
+        self.group = _get_group(group)
         self.add_prim_attr("no_eliminate", True)
 
     def infer_shape(self, x_shape):
@@ -1479,7 +1461,7 @@ class Receive(PrimitiveWithInfer):
             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/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1517,7 +1499,7 @@ class Receive(PrimitiveWithInfer):
         self.tag = sr_tag
         self.shape = shape
         self.dtype = dtype
-        self.group = group
+        self.group = _get_group(group)
         self.add_prim_attr("no_eliminate", True)
         valid_type = [mstype.float16, mstype.float32, mstype.float64, mstype.bfloat16,
                       mstype.int8, mstype.int16, mstype.int32, mstype.int64,
@@ -1695,6 +1677,8 @@ class _VirtualAssignAdd(PrimitiveWithInfer):
 
     def infer_dtype(self, x_dtype, y_dtype):
         return x_dtype
+
+
 virtual_assign_add = _VirtualAssignAdd()
 
 
@@ -1834,7 +1818,7 @@ class BatchISendIRecv(PrimitiveWithInfer):
             without any third-party or configuration file dependencies.
 
             Please see the `msrun start up
-            <https://www.mindspore.cn/tutorials/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1924,6 +1908,7 @@ class AlltoAllV(PrimitiveWithInfer):
         recv_numel_list(Union[tuple[int], list[int]]): split numel to gather from different remote rank.
         group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``, which
           means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
+        TODO:
 
     Inputs:
         - **input_x** (Tensor) - flatten tensor to scatter. The shape of tensor is :math:`(x_1)`.
@@ -1946,7 +1931,7 @@ class AlltoAllV(PrimitiveWithInfer):
             without any third-party or configuration file dependencies.
 
             Please see the `msrun start up
-            <https://www.mindspore.cn/tutorials/experts/zh-CN/master/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
@@ -1986,11 +1971,15 @@ class AlltoAllV(PrimitiveWithInfer):
     """
 
     @prim_attr_register
-    def __init__(self, send_numel_list, recv_numel_list, group=None):
+    def __init__(self, send_numel_list, recv_numel_list, group=None, split_sizes_empty=False):
         validator.check_value_type("send_numel_list", send_numel_list, [tuple, list], self.name)
         validator.check_value_type("recv_numel_list", recv_numel_list, [tuple, list], self.name)
-        if group is None:
-            group = GlobalComm.WORLD_COMM_GROUP
-        self.add_prim_attr('group', group)
+        self.group = GlobalComm.WORLD_COMM_GROUP if group is None else _get_group(group)
+        self.send_numel_list = send_numel_list
+        self.recv_numel_list = recv_numel_list
+        self.split_sizes_empty = split_sizes_empty
+        self.rank_size = get_group_size(self.group)
+
+        self.add_prim_attr('group', self.group)
         self.add_prim_attr('send_numel_list', send_numel_list)
         self.add_prim_attr('recv_numel_list', recv_numel_list)

mindspore/ops/operations/custom_ops.py

@@ -28,6 +28,8 @@ import subprocess
 import numpy as np
 import mindspore as ms
 from mindspore._c_expression import Oplib, typing
+from mindspore._c_expression import pyboost_custom_ext
+from mindspore.common._stub_tensor import _convert_stub
 from mindspore import context
 from mindspore.common import Tensor
 from mindspore.common import dtype as mstype
@@ -156,6 +158,55 @@ def _compile_aot(file):
     return func_path
 
 
+class _CustomExt(ops.PrimitiveWithInfer):
+    """
+    `Custom` primitive is used for PyBoost.
+    """
+
+    def __init__(self, func, out_shape=None, out_dtype=None, bprop=None):
+        super().__init__("CustomExt")
+        self.func = func
+        self.out_shape = out_shape
+        self.out_dtype = out_dtype
+        self.bprop = bprop
+
+    def __infer__(self, *args):
+        if callable(self.out_shape):
+            infer_shape = self.out_shape(*(x["shape"] for x in args))
+        else:
+            infer_shape = self.out_shape
+
+        if callable(self.out_dtype):
+            infer_dtype = self.out_dtype(*(x["dtype"] for x in args))
+        else:
+            infer_dtype = self.out_dtype
+
+        infer_value = None
+        if infer_shape is None:
+            logger.warning("'out_shape' is None. Add a placeholder instead. "
+                           "A CPP version of infer shape function is required "
+                           "in this case.")
+            infer_shape = (1,)
+        # after all automatic infer information fulfillment, throw error if infer_shape/infer_dtype is still None
+        if not isinstance(infer_shape, (tuple, list)):
+            raise TypeError("'out_shape' must be one of [tuple, list, function], but got {}".format(type(infer_shape)))
+
+        if not isinstance(infer_dtype, (typing.Type, tuple, list)):
+            raise TypeError("'out_dtype' must be one of [mindspore.dtype, tuple, list, function], but got {}"
+                            .format(type(infer_dtype)))
+
+        out = {
+            "shape": infer_shape,
+            "dtype": infer_dtype,
+            "value": infer_value,
+        }
+        return out
+
+    def get_bprop(self):
+        """return back propagation function"""
+        return self.bprop
+
+
 class Custom(ops.PrimitiveWithInfer):
     r"""
     `Custom` primitive is used for user defined operators and is to enhance the expressive ability of built-in
@@ -164,7 +215,7 @@ class Custom(ops.PrimitiveWithInfer):
     function if needed. Then these `Custom` objects can be directly used in neural networks.
     Detailed description and introduction of user-defined operators, including correct writing of parameters,
     please refer to `Custom Operators Tutorial
-    <https://www.mindspore.cn/tutorials/experts/en/master/operation/op_custom.html>`_ .
+    <https://www.mindspore.cn/docs/en/master/model_train/custom_program/op_custom.html>`_ .
 
     .. warning::
         - This is an experimental API that is subject to change.
@@ -174,7 +225,7 @@ class Custom(ops.PrimitiveWithInfer):
 
         - "hybrid": supports ["GPU", "CPU"].
         - "akg": supports ["GPU", "CPU"].
-        - "aot": supports ["GPU", "CPU", "ASCEDN"].
+        - "aot": supports ["GPU", "CPU", "Ascend"].
         - "pyfunc": supports ["CPU"].
         - "julia": supports ["CPU"].
 
@@ -249,20 +300,18 @@ class Custom(ops.PrimitiveWithInfer):
                        (ex. Custom(func="./reorganize.so:CustomReorganize", out_shape=[1], out_dtype=mstype.float32,
                        "aot"))
 
-                 b) ASCEND platform
-                 Before using Custom operators on the ASCEND platform, users must first develop custom operators
-                 based on Ascend C and compile them. For operator development, you can refer to the tutorial on
-                 `Quick Start for End-to-End Operator Development
-                 <https://www.hiascend.com/document/detail/zh/canncommercial/70RC1/operatordev/Ascendcopdevg/atlas_ascendc_10_0022.html>`_,
-                 and for compiling custom operators, you can use the `Offline Compilation of Ascend C Custom Operators
-                 <https://www.mindspore.cn/tutorials/experts/en/master/operation/op_custom_ascendc.html>` tool.
-                 When passing the operator's name into the func parameter, taking AddCustom as an example for the
-                 name given in the custom operator implementation, there are several ways to use it:
-
-                 - Usin TBE: func="AddCustom"
-                 - Using AclNN: func="aclnnAddCustom"
-                 - Inferring the shape of the operator through C++ derivation: func="infer_shape.cc:aclnnAddCustom",
-                   where infer_shape.cc is the shape derivation implemented in C++.
+                 b) Ascend platform.
+                 Before using Custom operators on the Ascend platform, users must first develop custom operators
+                 based on Ascend C and compile them. The complete development and usage process can refer to the
+                 tutorial `AOT-Type Custom Operators(Ascend) <https://www.mindspore.cn/docs/en/master/model_train/custom_program/operation/op_custom_ascendc.html>`_.
+                 By passing the name of the operator through the input parameter `func`, there are two usage methods
+                 based on the implementation of the infer shape function:
+
+                 - Python infer: If the operator's infer shape is implemented in Python, that is, the infer shape
+                   function is passed through the `out_shape` parameter, specify `func="CustomName"` .
+                 - C++ infer: If the operator's infer shape is implemented through C++, then pass the path of the
+                   infer shape implementation file in `func` and separate the operator name with `:`,
+                   for example: `func="add_custom_infer.cc:AddCustom"` .
 
               2. for "julia":
 
@@ -338,7 +387,7 @@ class Custom(ops.PrimitiveWithInfer):
             or the attributes of `func` differs in different targets.
 
     Supported Platforms:
-        ``GPU`` ``CPU`` ``ASCEND``
+        ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
         >>> import numpy as np
@@ -457,6 +506,12 @@ class Custom(ops.PrimitiveWithInfer):
 
         self.add_prim_attr("func_type", self.func_type)
         self._update_attr()
+        self.enable_pyboost = False
+        self.custom_pyboost = _CustomExt(self.func, self.out_shape, self.out_dtype, self.bprop)
+        if context.get_context("device_target") == "Ascend" and self.func_type == "aot":
+            self.enable_pyboost = True
+            for key, value in super().get_attr_dict().items():
+                self.custom_pyboost.add_prim_attr(key, value)
 
     def __infer__(self, *args):
         if callable(self.out_shape):
@@ -559,7 +614,7 @@ class Custom(ops.PrimitiveWithInfer):
             raise TypeError(
                 "{}, 'func' should be like 'file_name:func_name', but got {}".format(
                     self.log_prefix, self.func))
-        file_path = os.path.abspath(file_name_list[0])
+        file_path = os.path.realpath(file_name_list[0])
         if os.environ.get('MS_CUSTOM_AOT_WHITE_LIST') is None:
             if Custom.custom_aot_warning:
                 logger.info("{}, no white list is set and it might cause problems. "
@@ -567,7 +622,7 @@ class Custom(ops.PrimitiveWithInfer):
                             .format(self.log_prefix))
                 Custom.custom_aot_warning = False
         else:
-            legal_path = os.path.abspath(os.environ.get('MS_CUSTOM_AOT_WHITE_LIST'))
+            legal_path = os.path.realpath(os.environ.get('MS_CUSTOM_AOT_WHITE_LIST'))
             if legal_path not in file_path:
                 raise TypeError(
                     "{}, the legal path for the file is {}, but the file is {}".format(
@@ -1063,3 +1118,12 @@ class Custom(ops.PrimitiveWithInfer):
         infer_value = Tensor(fake_output) if enable_infer_value else None
 
         return infer_shape, infer_dtype, infer_value
+
+    def __call__(self, *args):
+        if self.enable_pyboost:
+            return _convert_stub(pyboost_custom_ext(self.custom_pyboost, [args]))
+        should_elim, output = self.check_elim(*args)
+        if should_elim:
+            return output
+        # pylint: disable=protected-access
+        return ops.primitive._run_op(self, self.name, args)

mindspore/ops/operations/debug_ops.py

@@ -15,17 +15,17 @@
 """debug_ops"""
 import os
 import stat
-from types import FunctionType, MethodType
 
 import numpy as np
 from mindspore import log as logger
-from mindspore._c_expression import security
+from mindspore._c_expression import security, HookType
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 from mindspore.ops.primitive import prim_attr_register, Primitive, PrimitiveWithInfer
+from mindspore._checkparam import check_hook_fn
 
 
 SUMMARY_TENSOR_CACHE = []
@@ -64,6 +64,8 @@ class ScalarSummary(Primitive):
     which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
     mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -122,6 +124,8 @@ class ImageSummary(Primitive):
     SummaryRecord or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
     mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable, it must not be an empty string.
@@ -173,6 +177,8 @@ class TensorSummary(Primitive):
     or SummaryCollector, which specify the directory of the summary file. The summary file can
     be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
     mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable.
@@ -228,9 +234,6 @@ class TensorDump(Primitive):
     """
     Save the Tensor as an npy file in numpy format.
 
-    The file name will automatically have a prefix added based on the execution order. For example, if `file` is `a`,
-    the first saved file will be named `0_a.npy`, and the second one will be named `1_a.npy`, and so on.
-
     .. warning::
         - If a large amount of data is stored within a short period, it may lead to memory overflow on the device side.
           Consider slicing the data to reduce the data scale.
@@ -238,6 +241,34 @@ class TensorDump(Primitive):
           too quickly, data loss may occur. You need to actively control the destruction time of the main process,
           such as using sleep.
 
+    Args:
+        input_output (str, optional): Used to control Tensordump behavior.
+            Available value is one of ['in', 'out', 'all']. Default value is ``out``.
+
+            In case of OpA --> RedistributionOps --> OpB,
+            The dump data of OpA's output is not equal to OpB's input (Due to the redistribution operators).
+            So the parameter input_output is to handle this situation.
+
+            Assuming OpA's output is used as both Tensordump's input parameter and OpB's input parameter.
+            Different requirements of saving dump data can be achieved by configuring parameter input_output:
+
+            - If the input_output is 'out', the dump data contains only OpA's output slice.
+            - If the input_output is 'all', the dump data contains both OpA's output slice and OpB's input slice.
+            - If the input_output is 'in', the dump data contains only OpB's input slice.
+
+            For input_output is 'all' or 'in', the input slice npy file format is:
+            id_fileName_cNodeID_dumpMode_rankID.npy.
+
+            For input_output is 'out' or 'all' the output slice npy file format is:
+            id_fileName.npy.
+
+            - id: An auto increment ID.
+            - fileName: Value of the parameter file
+              (if parameter file_name is a user-specified path, the value of fileName is the last level of the path).
+            - cNodeID: The node ID of the Tensordump node in the step_parallel_end.ir file.
+            - dumpMode: Value of the parameter input_output.
+            - rankID: Logical device id.
+
     Inputs:
         - **file** (str) - The path of the file to be saved.
         - **input_x** (Tensor) - Input Tensor of any dimension.
@@ -280,7 +311,7 @@ class TensorDump(Primitive):
          [6. 7. 8. 9.]]
     """
     @prim_attr_register
-    def __init__(self):
+    def __init__(self, input_output='out'):
         """Initialize TensorDump."""
         if security.enable_security():
             raise ValueError('The TensorDump is not supported, please without `-s on` and recompile source.')
@@ -314,6 +345,8 @@ class HistogramSummary(Primitive):
     It must be used with SummaryRecord or SummaryCollector, which specify the directory of the summary file.
     The summary file can be loaded and shown by MindInsight, see `MindInsight documents <https://www.mindspore.cn/
     mindinsight/docs/en/master/index.html>`_ for details.
+    In Ascend platform with graph mode, can set environment variables `MS_DUMP_SLICE_SIZE` and `MS_DUMP_WAIT_TIME`
+    to solve operator execution failure when calling this operator intensively.
 
     Inputs:
         - **name** (str) - The name of the input variable.
@@ -499,16 +532,15 @@ class HookBackward(PrimitiveWithInfer):
     def __init__(self, hook_fn, cell_id=""):
         """Initialize HookBackward."""
         super(HookBackward, self).__init__(self.__class__.__name__)
-        if not isinstance(hook_fn, (FunctionType, MethodType)):
-            raise TypeError(f"For '{self.name}', the type of 'hook_fn' must be python function, "
-                            f"but got {type(hook_fn)}.")
+        if not check_hook_fn("HookBackward", hook_fn):
+            return
         if cell_id != "":
             logger.warning(f"The args 'cell_id' of HookBackward will be removed in a future version. If the value of "
                            f"'cell_id' is set, the hook function will not work.")
         self.add_prim_attr("cell_id", cell_id)
         self.init_attrs["cell_id"] = cell_id
         self.cell_id = cell_id
-        self.add_backward_hook_fn(hook_fn)
+        self.set_hook_fn(hook_fn, HookType.HookBackward)
 
     def infer_shape(self, *inputs_shape):
         if len(inputs_shape) == 1:

mindspore/ops/operations/manually_defined/_inner.py

@@ -23,6 +23,7 @@ from mindspore.common import dtype as mstype
 from mindspore import _checkparam as validator
 from mindspore.common._decorator import deprecated
 from mindspore.ops.primitive import prim_attr_register, Primitive
+from mindspore import log as logger
 
 
 class ScalarCast(Primitive):
@@ -59,3 +60,14 @@ class ScalarCast(Primitive):
         value = np.cast[np_dtype.lower()](input_x)
         value = value.item()
         return value
+
+
+class TensorReport(Primitive):
+    @prim_attr_register
+    def __init__(self):
+        """Initialize TensorReport"""
+        self.add_prim_attr("side_effect_io", True)
+        self.add_prim_attr("channel_name", "ms_tensor_report")
+
+    def __call__(self, file, input_x):
+        logger.warning("TensorReport doesn't support pynative mode.")