mindspore 2.6.0rc1__cp311-cp311-win_amd64.whl → 2.7.0__cp311-cp311-win_amd64.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 = []
 
 
@@ -310,6 +309,8 @@ class TensorDump(Primitive):
         """Initialize TensorDump."""
         if security.enable_security():
             raise ValueError('The TensorDump is not supported, please without `-s on` and recompile source.')
+        if input_output not in ['in', 'out']:
+            raise ValueError(f"The 'input_output' argument should be one of ['in', 'out'], but got: {input_output}")
         self.add_prim_attr("side_effect_io", True)
         self.add_prim_attr("channel_name", "ms_tensor_dump")
 
@@ -396,6 +397,10 @@ class InsertGradientOf(Primitive):
     """
     Attaches callback to the graph node that will be invoked on the node's gradient.
 
+    .. warning::
+        In the callback, exercise caution when using side-effect operators,
+        such as the TensorDump operator, as current support is incomplete.
+
     Args:
         f (Function): MindSpore's Function. Callback function.
 
@@ -459,15 +464,65 @@ 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 __call__(self, path, x, input_output):
+        def _dump_hook(dout):
+            P.TensorDump()(path, dout)
+            return dout
+        x = P.InsertGradientOf(_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
@@ -978,7 +977,7 @@ class ScalarToTensor(PrimitiveWithInfer):
     def __call__(self, x, dtype=mstype.float32):
         validator.check_value_type("x", x, [bool, int, float], self.name)
         validator.check_subclass("dtype", dtype, mstype.number, self.name)
-        data_type = mstype.dtype_to_nptype(dtype)
+        data_type = mstype._dtype_to_nptype(dtype)  # pylint:disable=protected-access
         return Tensor(np.array(x, data_type), dtype=dtype)
 
 
@@ -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):
@@ -1151,7 +1149,7 @@ def scalar_cast(input_x, input_y):
     Args:
         input_x (scalar): The input scalar.
         input_y (mindspore.dtype): The type to be cast. Only constant value is allowed.
-            The value should only be mindspore.int64, mindspore.float64, or mindspore.bool\_.
+            The value should only be mindspore.int64, mindspore.float64, or mindspore.bool.
 
     Returns:
         Scalar, the type is the same as the python type corresponding to `input_y`.
@@ -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):
@@ -1716,7 +1714,7 @@ def infer_value_for_Arange(start, end, step, dtype=None):
         if has_float:
             np_dtype = np.float32
     else:
-        np_dtype = mstype.dtype_to_nptype(typing.type_id_to_type(dtype))
+        np_dtype = mstype._dtype_to_nptype(typing.type_id_to_type(dtype))  # pylint:disable=protected-access
     return Tensor(np.arange(start, end, step, dtype=np_dtype))
 
 
@@ -1740,7 +1738,7 @@ def _infer_value_for_ReduceExtand(input_x, axis, keep_dims, dtype, prim_name):
             else:
                 axis = tuple(range(len(value.shape)))
             if dtype is not None:
-                np_dtype = mstype.dtype_to_nptype(typing.type_id_to_type(dtype))
+                np_dtype = mstype._dtype_to_nptype(typing.type_id_to_type(dtype))  # pylint:disable=protected-access
                 value = np_reduce_extand_func(value, axis, dtype=np_dtype, keepdims=keep_dims)
             else:
                 value = np_reduce_extand_func(value, axis, keepdims=keep_dims)
@@ -1773,7 +1771,7 @@ def infer_value_for_Cast(x, dst_type_enum=None):
     if x is None or dst_type_enum is None:
         return None
     dst_type = typing.type_id_to_type(dst_type_enum)
-    src_type = mstype.get_py_obj_dtype(x)
+    src_type = mstype._get_py_obj_dtype(x)  # pylint:disable=protected-access
     validator.check_subclass("input_x", src_type, [mstype.tensor_type, mstype.number], "Cast")
     validator.check_subclass("type", dst_type, mstype.number, "Cast")
 
@@ -1783,7 +1781,7 @@ def infer_value_for_Cast(x, dst_type_enum=None):
         dst_type = dst_type.element_type()
 
     value = None
-    np_dst_type = mstype.dtype_to_nptype(dst_type)
+    np_dst_type = mstype._dtype_to_nptype(dst_type)  # pylint:disable=protected-access
     if isinstance(x, (int, float)):
         value = Tensor(np.array(x).astype(np_dst_type), dtype=dst_type)
     else:
@@ -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):
@@ -2840,8 +2839,8 @@ class WhileLoop(Primitive):
             while cond_func(val):
                 val = loop_func(val)
         except Exception as e:
-            raise ValueError("Invalid loop_func, please check input arguments and \
-                             return value, error info: {}".format(e))
+            raise ValueError(f"Invalid loop_func, please check input arguments and "
+                             f"return value, error info: {e}")
         return val
 
 
@@ -2936,8 +2935,8 @@ class Scan(Primitive):
                 ys.append(y)
                 i = i + 1
         except Exception as e:
-            raise ValueError("Invalid loop_func, please check input arguments and \
-                             return value, error info: {}".format(e))
+            raise ValueError(f"Invalid loop_func, please check input arguments and "
+                             f"return value, error info: {e}")
         return carry, ys
 
 
@@ -3012,6 +3011,6 @@ class ForiLoop(Primitive):
             for i in range(lower, upper):
                 val = loop_func(i, val)
         except Exception as e:
-            raise ValueError("Invalid loop_func, please check input arguments and \
-                             return value, error info: {}".format(e))
+            raise ValueError(f"Invalid loop_func, please check input arguments and "
+                             f"return value, error info: {e}")
         return val

mindspore/ops/operations/math_ops.py

@@ -882,7 +882,7 @@ class Sub(_MathBinaryOp):
     Note:
         - When the two inputs have different shapes, they must be able to broadcast to a common shape.
         - The two inputs can not be bool type at the same time,
-          [True, Tensor(True, bool\_), Tensor(np.array([True]), bool\_)] are all considered bool type.
+          [True, Tensor(True), Tensor(np.array([True]))] are all considered bool type.
         - The two inputs comply with the implicit type conversion rules to make the data types
           consistent.
 
@@ -890,7 +890,7 @@ class Sub(_MathBinaryOp):
         - **x** (Union[Tensor, number.Number, bool]) - The first input is a number.Number or
           a bool or a tensor whose data type is
           `number <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.dtype.html>`_ or
-          `bool_ <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.dtype.html>`_.
+          `bool <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.dtype.html>`_.
         - **y** (Union[Tensor, number.Number, bool]) - The second input, when the first input is a Tensor,
           the second input should be a number.Number or bool value, or a Tensor whose data type is number or bool.
 
@@ -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.
@@ -1289,10 +1289,10 @@ class DivNoNan(Primitive):
         - **x1** (Union[Tensor, number.Number, bool]) - The first input is a number.Number or
           a bool or a tensor whose data type is
           `number <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.dtype.html>`_ or
-          `bool_ <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.dtype.html>`_.
+          `bool <https://www.mindspore.cn/docs/en/master/api_python/mindspore/mindspore.dtype.html>`_.
         - **x2** (Union[Tensor, number.Number, bool]) - The second input is a number.Number or
-          a bool when the first input is a bool or a tensor whose data type is number or bool\_.
-          When the first input is Scalar, the second input must be a Tensor whose data type is number or bool\_.
+          a bool when the first input is a bool or a tensor whose data type is number or bool.
+          When the first input is Scalar, the second input must be a Tensor whose data type is number or bool.
 
     Outputs:
         Tensor, the shape is the same as the one after broadcasting,
@@ -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,
@@ -42,7 +41,7 @@ from ..auto_generate import (CeLU, Flatten, LogSoftmax, LogSoftmaxExt, GLU, ReLU
                              UpsampleNearest3D, UpsampleTrilinear3D,
                              SoftMarginLoss, UpsampleBilinear2D, UpsampleLinear1D,
                              BinaryCrossEntropy, BCEWithLogitsLoss, SoftShrink, AdaptiveMaxPool2D,
-                             SmoothL1Loss)
+                             SmoothL1Loss, KvScaleCache)
 from .manually_defined import BatchNorm
 
 
@@ -6434,39 +6433,8 @@ class Conv3D(Primitive):
 
     Outputs:
         Tensor, the value that applied 3D convolution. The shape is :math:`(N, C_{out}, D_{out}, H_{out}, W_{out})`.
-
-        `pad_mode` is ``"same"``:
-
-        .. math::
-            \begin{array}{ll} \\
-                D_{out} = \left \lceil{\frac{D_{in}}{\text{stride[0]}}} \right \rceil \\
-                H_{out} = \left \lceil{\frac{H_{in}}{\text{stride[1]}}} \right \rceil \\
-                W_{out} = \left \lceil{\frac{W_{in}}{\text{stride[2]}}} \right \rceil \\
-            \end{array}
-
-        `pad_mode` is ``"valid"``:
-
-        .. math::
-            \begin{array}{ll} \\
-                D_{out} = \left \lfloor{\frac{D_{in} - \text{dilation[0]} \times (\text{kernel_size[0]} - 1) }
-                {\text{stride[0]}} + 1} \right \rfloor \\
-                H_{out} = \left \lfloor{\frac{H_{in} - \text{dilation[1]} \times (\text{kernel_size[1]} - 1) }
-                {\text{stride[1]}} + 1} \right \rfloor \\
-                W_{out} = \left \lfloor{\frac{W_{in} - \text{dilation[2]} \times (\text{kernel_size[2]} - 1) }
-                {\text{stride[2]}} + 1} \right \rfloor \\
-            \end{array}
-
-        `pad_mode` is ``"pad"``:
-
-        .. math::
-            \begin{array}{ll} \\
-                D_{out} = \left \lfloor{\frac{D_{in} + pad[0] + pad[1] - (\text{dilation[0]} - 1) \times
-                \text{kernel_size[0]} - 1 }{\text{stride[0]}} + 1} \right \rfloor \\
-                H_{out} = \left \lfloor{\frac{H_{in} + pad[2] + pad[3] - (\text{dilation[1]} - 1) \times
-                \text{kernel_size[1]} - 1 }{\text{stride[1]}} + 1} \right \rfloor \\
-                W_{out} = \left \lfloor{\frac{W_{in} + pad[4] + pad[5] - (\text{dilation[2]} - 1) \times
-                \text{kernel_size[2]} - 1 }{\text{stride[2]}} + 1} \right \rfloor \\
-            \end{array}
+        To see how different pad modes affect the output shape, please refer to
+        :class:`mindspore.nn.Conv3d` for more details.
 
     Raises:
         TypeError: If `out_channel` or `group` is not an int.
@@ -7151,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) "
@@ -7402,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__ = (
@@ -9289,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)
 
@@ -461,7 +458,7 @@ class Primitive(Primitive_):
 
             - If the computation involves something like randomization or global variable, the equivalence
               is not guaranteed currently.
-            - Not supported in pynative mode
+            - Should only be used in Graph mode or in gradient functions that are decorated by @jit.
 
         Args:
             mode (bool): Specifies whether the primitive is recomputed. Default: ``True`` .
@@ -469,7 +466,7 @@ class Primitive(Primitive_):
         Examples:
             >>> import numpy as np
             >>> import mindspore as ms
-            >>> from mindspore import Tensor, ops, nn
+            >>> from mindspore import Tensor, ops, nn, jit
             >>> class NetRecompute(nn.Cell):
             ...     def __init__(self):
             ...         super(NetRecompute,self).__init__()
@@ -484,6 +481,7 @@ class Primitive(Primitive_):
             ...         super(GradNet,self).__init__()
             ...         self.network = network
             ...         self.grad = ops.GradOperation()
+            ...     @jit
             ...     def construct(self, x):
             ...         g_out = self.grad(self.network)(x)
             ...         return g_out
@@ -495,8 +493,6 @@ class Primitive(Primitive_):
             >>> print(a)
             [0. 0.5]
         """
-        if context.get_context("mode") == context.PYNATIVE_MODE:
-            raise TypeError("Recompute is not supported in pynative mode currently.")
         Validator.check_bool(mode)
         self.add_prim_attr("recompute", mode)
         return self
@@ -513,8 +509,6 @@ class Primitive(Primitive_):
         Args:
             backward_prefetch(Union[str, int]): Specifies whether the activation is prefetched in backward pass.
         """
-        if context.get_context("mode") == context.PYNATIVE_MODE:
-            raise ValueError("Offload is not supported in pynative mode currently.")
         self.add_prim_attr("offload", True)
         if isinstance(backward_prefetch, str):
             Validator.check_string(backward_prefetch, ['Auto'], 'backward_prefetch', 'Primitive._offload')
@@ -554,10 +548,6 @@ class Primitive(Primitive_):
         Validator.check_non_negative_int(rank_id, "rank_id", "Primitive.place")
         Validator.check_string(role, "MS_WORKER", "role", "Primitive.place")
 
-        if context.get_context("mode") == context.PYNATIVE_MODE:
-            raise RuntimeError("You are calling Primitive.place in pynative mode."
-                               "It's only supported in graph mode. Please switch to graph mode.")
-
         # Get the execution context and check whether calling of this 'place' method is valid.
         # This is because placing operators to arbitrary processes while other distributed training mode
         # is enabled is very unpredictable and may cause fatal error.
@@ -1032,7 +1022,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