mindspore 2.4.0__cp311-cp311-manylinux1_x86_64.whl → 2.4.1__cp311-cp311-manylinux1_x86_64.whl

mindspore/ops/auto_generate/gen_ops_prim.py

@@ -81,6 +81,7 @@ from mindspore._c_expression import pyboost_copy_ext
 from mindspore._c_expression import pyboost_copy
 from mindspore._c_expression import pyboost_cos
 from mindspore._c_expression import pyboost_cosh
+from mindspore._c_expression import pyboost_count_nonzero
 from mindspore._c_expression import pyboost_cross
 from mindspore._c_expression import pyboost_cummax
 from mindspore._c_expression import pyboost_cummin_ext
@@ -116,7 +117,6 @@ from mindspore._c_expression import pyboost_gather_d_grad_v2
 from mindspore._c_expression import pyboost_gather_d
 from mindspore._c_expression import pyboost_gelu_grad
 from mindspore._c_expression import pyboost_gelu
-from mindspore._c_expression import pyboost_generator
 from mindspore._c_expression import pyboost_greater_equal
 from mindspore._c_expression import pyboost_greater
 from mindspore._c_expression import pyboost_grid_sampler_2d_grad
@@ -195,6 +195,7 @@ from mindspore._c_expression import pyboost_muls
 from mindspore._c_expression import pyboost_multinomial_ext
 from mindspore._c_expression import pyboost_mv
 from mindspore._c_expression import pyboost_nan_to_num
+from mindspore._c_expression import pyboost_ne_scalar
 from mindspore._c_expression import pyboost_neg
 from mindspore._c_expression import pyboost_non_zero_ext
 from mindspore._c_expression import pyboost_non_zero
@@ -214,6 +215,11 @@ from mindspore._c_expression import pyboost_prelu
 from mindspore._c_expression import pyboost_prod_ext
 from mindspore._c_expression import pyboost_rand_ext
 from mindspore._c_expression import pyboost_rand_like_ext
+from mindspore._c_expression import pyboost_randint_like
+from mindspore._c_expression import pyboost_randint
+from mindspore._c_expression import pyboost_randn_like
+from mindspore._c_expression import pyboost_randn
+from mindspore._c_expression import pyboost_randperm_ext
 from mindspore._c_expression import pyboost_reciprocal
 from mindspore._c_expression import pyboost_reduce_all
 from mindspore._c_expression import pyboost_reduce_any
@@ -250,6 +256,7 @@ from mindspore._c_expression import pyboost_scatter_add_ext
 from mindspore._c_expression import pyboost_scatter
 from mindspore._c_expression import pyboost_scatter_value
 from mindspore._c_expression import pyboost_searchsorted
+from mindspore._c_expression import pyboost_select_ext
 from mindspore._c_expression import pyboost_select
 from mindspore._c_expression import pyboost_select_v2
 from mindspore._c_expression import pyboost_selu_ext
@@ -317,6 +324,7 @@ from mindspore._c_expression import pyboost_zeros
 from mindspore._c_expression import pyboost_add_rmsnorm_quant_v2
 from mindspore._c_expression import pyboost_dynamic_quant_ext
 from mindspore._c_expression import pyboost_grouped_matmul
+from mindspore._c_expression import pyboost_kv_cache_scatter_update
 from mindspore._c_expression import pyboost_moe_finalize_routing
 from mindspore._c_expression import pyboost_quant_batch_matmul
 from mindspore._c_expression import pyboost_quant_v2
@@ -3770,6 +3778,36 @@ class Cosh(Primitive):
 cosh_op=Cosh()
 
 
+class CountNonZero(Primitive):
+    r"""
+    .. code-block::
+        
+        prim = ops.CountNonZero()
+        out = prim(input, dim)
+        
+    is equivalent to
+    
+    .. code-block::
+    
+        ops.count_nonzero(input, dim)
+        
+    Refer to :func:`mindspore.ops.count_nonzero` for more details.
+    """
+    __mindspore_signature__ =  (
+        sig.make_sig('input'),
+        sig.make_sig('dim', default=None),
+    )
+
+    @prim_arg_register
+    def __init__(self):
+        pass
+
+    def __call__(self, input, dim=None):
+          return _convert_stub(pyboost_count_nonzero(self, [input, dim]))
+
+count_nonzero_op=CountNonZero()
+
+
 class Cross(Primitive):
     r"""
     Returns the cross product of vectors in dimension `dim` of input and other.
@@ -6490,7 +6528,8 @@ class Generator(Primitive):
         self.add_prim_attr("side_effect_mem", True)
 
     def __call__(self, cmd, inputs):
-          return _convert_stub(pyboost_generator(self, [cmd, inputs]))
+          return super().__call__(cmd, inputs)
+
 
 generator_op=Generator()
 
@@ -8678,6 +8717,9 @@ class LinSpaceExt(Primitive):
         &output = [start, start+step, start+2*step, ... , end]
         \end{aligned}
 
+    .. warning::
+        Atlas training series does not support int16 dtype currently.
+
     Inputs:
         - **start** (Union[float, int]) - Start value of interval.
           It can be a float or integer.
@@ -11245,6 +11287,115 @@ class RandLikeExt(Primitive):
 rand_like_ext_op=RandLikeExt()
 
 
+class RandIntLike(Primitive):
+    r"""
+    
+    """
+    __mindspore_signature__ =  (
+        sig.make_sig('input'),
+        sig.make_sig('low'),
+        sig.make_sig('high'),
+        sig.make_sig('seed'),
+        sig.make_sig('offset'),
+        sig.make_sig('dtype', default=None),
+    )
+
+    @prim_arg_register
+    def __init__(self):
+        pass
+
+    def __call__(self, input, low, high, seed, offset, dtype=None):
+          return _convert_stub(pyboost_randint_like(self, [input, low, high, seed, offset, dtype if dtype is None else dtype_to_type_id('RandIntLike', 'dtype', dtype)]))
+
+randint_like_op=RandIntLike()
+
+
+class RandInt(Primitive):
+    r"""
+    
+    """
+    __mindspore_signature__ =  (
+        sig.make_sig('low'),
+        sig.make_sig('high'),
+        sig.make_sig('shape'),
+        sig.make_sig('seed'),
+        sig.make_sig('offset'),
+        sig.make_sig('dtype', default=None),
+    )
+
+    @prim_arg_register
+    def __init__(self):
+        pass
+
+    def __call__(self, low, high, shape, seed, offset, dtype=None):
+          return _convert_stub(pyboost_randint(self, [low, high, shape, seed, offset, dtype if dtype is None else dtype_to_type_id('RandInt', 'dtype', dtype)]))
+
+randint_op=RandInt()
+
+
+class RandnLike(Primitive):
+    r"""
+    
+    """
+    __mindspore_signature__ =  (
+        sig.make_sig('input'),
+        sig.make_sig('seed'),
+        sig.make_sig('offset'),
+        sig.make_sig('dtype', default=None),
+    )
+
+    @prim_arg_register
+    def __init__(self):
+        pass
+
+    def __call__(self, input, seed, offset, dtype=None):
+          return _convert_stub(pyboost_randn_like(self, [input, seed, offset, dtype if dtype is None else dtype_to_type_id('RandnLike', 'dtype', dtype)]))
+
+randn_like_op=RandnLike()
+
+
+class Randn(Primitive):
+    r"""
+    
+    """
+    __mindspore_signature__ =  (
+        sig.make_sig('shape'),
+        sig.make_sig('seed'),
+        sig.make_sig('offset'),
+        sig.make_sig('dtype', default=None),
+    )
+
+    @prim_arg_register
+    def __init__(self):
+        pass
+
+    def __call__(self, shape, seed, offset, dtype=None):
+          return _convert_stub(pyboost_randn(self, [shape, seed, offset, dtype if dtype is None else dtype_to_type_id('Randn', 'dtype', dtype)]))
+
+randn_op=Randn()
+
+
+class RandpermExt(Primitive):
+    r"""
+    
+    """
+    __mindspore_signature__ =  (
+        sig.make_sig('n'),
+        sig.make_sig('seed'),
+        sig.make_sig('offset'),
+        sig.make_sig('dtype', default=mstype.int64),
+    )
+
+    @prim_arg_register
+    def __init__(self):
+        pass
+
+    def __call__(self, n, seed, offset, dtype=mstype.int64):
+          return _convert_stub(pyboost_randperm_ext(self, [n, seed, offset, dtype_to_type_id('RandpermExt', 'dtype', dtype)]))
+
+randperm_ext_op=RandpermExt()
+
+
 class RandpermV2(Primitive):
     r"""
     .. code-block::
@@ -13631,8 +13782,7 @@ class SelectExt(Primitive):
         pass
 
     def __call__(self, input, dim, index):
-          return super().__call__(input, dim, index)
-
+          return _convert_stub(pyboost_select_ext(self, [input, dim, index]))
 
 select_ext_op=SelectExt()
 
@@ -16494,8 +16644,7 @@ class KVCacheScatterUpdate(Primitive):
         self.add_prim_attr("side_effect_mem", True)
 
     def __call__(self, var, indices, updates, axis, reduce='none'):
-          return super().__call__(var, indices, updates, axis, str_to_enum('KVCacheScatterUpdate', 'reduce', reduce))
-
+          return _convert_stub(pyboost_kv_cache_scatter_update(self, [var, indices, updates, axis, str_to_enum('KVCacheScatterUpdate', 'reduce', reduce)]))
 
 kv_cache_scatter_update_op=KVCacheScatterUpdate()
 

mindspore/ops/function/array_func.py

@@ -1312,6 +1312,8 @@ def unique_with_pad(x, pad_num):
 
     .. warning::
         :func:`mindspore.ops.unique_with_pad` is deprecated from version 2.4 and will be removed in a future version.
+        Please use the :func:`mindspore.ops.unique` combined with :func:`mindspore.ops.pad` to realize
+        the same function.
 
     Args:
         x (Tensor): The tensor need to be unique. Must be 1-D vector with types: int32, int64.

mindspore/ops/function/math_func.py

@@ -2761,6 +2761,9 @@ def linspace_ext(start, end, steps, *, dtype=None):
         &output = [start, start+step, start+2*step, ... , end]
         \end{aligned}
 
+    .. warning::
+        Atlas training series does not support int16 dtype currently.
+
     Args:
         start (Union[float, int]): Start value of interval.
             It can be a float or integer.
@@ -7518,7 +7521,7 @@ def norm_ext(input, p='fro', dim=None, keepdim=False, *, dtype=None):
         This is an experimental API that is subject to change or deletion.
 
     Args:
-        input (Tensor): The input of LogSigmoid with data type of bfloat16, float16 or float32.
+        input (Tensor): The input of norm with data type of bfloat16, float16 or float32.
             The shape is :math:`(*)` where :math:`*` means, any number of additional dimensions.
         p (Union[int, float, inf, -inf, 'fro', 'nuc'], optional): norm's mode. refer to the table above for
             behavior. Default: ``fro`` .
@@ -7554,6 +7557,9 @@ def norm_ext(input, p='fro', dim=None, keepdim=False, *, dtype=None):
         >>> print(ops.function.math_func.norm_ext(x, 2.0))
         38.327538
     """
+    if not isinstance(input, (Tensor, Tensor_)):
+        raise TypeError(f"For `norm_ext`, the `input` must be Tensor!, but get {type(input)}.")
+
     if (dim is not None) or keepdim or (dtype is not None):
         raise ValueError(f"For `norm_ext`, the value of `dim`, `keepdim` and `dtype` must be default value currently.")
 

mindspore/ops/function/random_func.py

@@ -30,7 +30,8 @@ from mindspore.common.api import _function_forbid_reuse
 from mindspore.ops.auto_generate import randperm
 from mindspore.common.generator import default_generator
 from mindspore.ops.auto_generate import UniformExt, NormalTensorTensor, \
-    NormalTensorFloat, NormalFloatTensor, NormalFloatFloat, RandExt, RandLikeExt, MultinomialExt
+    NormalTensorFloat, NormalFloatTensor, NormalFloatFloat, RandExt, RandLikeExt, MultinomialExt, \
+    Randn, RandnLike, RandInt, RandIntLike, RandpermExt
 
 normal_tensor_tensor_op = NormalTensorTensor()
 normal_tensor_float_op = NormalTensorFloat()
@@ -42,10 +43,15 @@ real_div_ = P.RealDiv()
 reshape_ = P.Reshape()
 shape_ = P.Shape()
 top_k_ = P.TopK()
+randperm_ext_ = RandpermExt()
 uniform_ = UniformExt()
 rand_ext_ = RandExt()
 rand_like_ext_ = RandLikeExt()
 multinomial_ext_ = MultinomialExt()
+randn_ = Randn()
+randn_like_ = RandnLike()
+randint_ = RandInt()
+randint_like_ = RandIntLike()
 generator_step_ = Tensor(10, mstype.int64)
 
 
@@ -287,7 +293,8 @@ def uniform_ext(tensor, a, b, generator=None):
     """
     if generator is None:
         generator = default_generator
-    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
     return uniform_(tensor, a, b, seed, offset)
 
 
@@ -755,7 +762,8 @@ def normal_ext(mean=0.0, std=1.0, size=None, generator=None):
     """
     if generator is None:
         generator = default_generator
-    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
 
     is_mean_tensor = isinstance(mean, Tensor)
     is_std_tensor = isinstance(std, Tensor)
@@ -1129,7 +1137,8 @@ def rand_ext(*size, generator=None, dtype=None):
     """
     if not generator:
         generator = default_generator
-    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
     return rand_ext_(size, seed, offset, dtype)
 
 
@@ -1163,10 +1172,174 @@ def rand_like_ext(input, *, dtype=None):
         >>> print(ops.function.random_func.rand_like_ext(a, dtype=ms.float32).shape)
         (2, 3)
     """
-    seed, offset = default_generator._step(generator_step_)  # pylint: disable=protected-access
+    seed, offset = default_generator._step(  # pylint: disable=protected-access
+        generator_step_)
     return rand_like_ext_(input, seed, offset, dtype)
 
 
+@_function_forbid_reuse
+def randn_ext(*size, generator=None, dtype=None):
+    r"""
+    Returns a new tensor filled with numbers from the normal distribution over an interval :math:`[0, 1)`
+    based on the given shape and dtype.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        size (Union[int, tuple(int), list(int)]): Shape of the new tensor, e.g. :math:`(2, 3)` or :math:`2`.
+
+    Keyword Args:
+        generator (:class:`mindspore.Generator`, optional): a pseudorandom number generator.
+            Default: ``None``, uses the default pseudorandom number generator.
+        dtype (:class:`mindspore.dtype`, optional): Designated tensor dtype, it must be float type. If None,
+            `mindspore.float32` will be applied. Default: ``None`` .
+
+    Returns:
+        Tensor, with the designated shape and dtype, filled with random numbers from the normal distribution on
+        the interval :math:`[0, 1)`.
+
+    Raises:
+        ValueError: If `dtype` is not a `mstype.float_type` type.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore import ops
+        >>> print(ops.function.random_func.randn_ext(2, 3).shape)
+        (2, 3)
+    """
+    if not generator:
+        generator = default_generator
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
+    return randn_(size, seed, offset, dtype)
+
+
+@_function_forbid_reuse
+def randn_like_ext(input, *, dtype=None):
+    r"""
+    Returns a new tensor filled with numbers from the normal distribution over an interval :math:`[0, 1)`
+    based on the given dtype and shape of the input tensor.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        input (Tensor): Input Tensor to specify the output shape and its default dtype.
+
+    Keyword Args:
+        dtype (:class:`mindspore.dtype`, optional): Designated tensor dtype, it must be float type. If None,
+            the same dtype of `input` will be applied. Default: ``None`` .
+
+    Returns:
+        Tensor, with the designated shape and dtype, filled with random numbers from the normal distribution on
+        the interval :math:`[0, 1)`.
+
+    Raises:
+        ValueError: If `dtype` is not a `mstype.float_type` type.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import Tensor, ops
+        >>> a = Tensor([[2, 3, 4], [1, 2, 3]])
+        >>> print(ops.function.random_func.randn_like_ext(a, dtype=ms.float32).shape)
+        (2, 3)
+    """
+    seed, offset = default_generator._step(  # pylint: disable=protected-access
+        generator_step_)
+    return randn_like_(input, seed, offset, dtype)
+
+
+@_function_forbid_reuse
+def randint_ext(low, high, size, *, generator=None, dtype=None):
+    r"""
+    Returns a new tensor filled with integer numbers from the uniform distribution over an interval :math:`[low, high)`
+    based on the given shape and dtype.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        low (int): the lower bound of the generated random number
+        high (int): the upper bound of the generated random number
+        size (Union[tuple(int), list(int)]): Shape of the new tensor, e.g. :math:`(2, 3)`.
+
+    Keyword Args:
+        generator (:class:`mindspore.Generator`, optional): a pseudorandom number generator.
+            Default: ``None``, uses the default pseudorandom number generator.
+        dtype (:class:`mindspore.dtype`, optional): Designated tensor dtype. If None,
+            `mindspore.int64` will be applied. Default: ``None`` .
+
+    Returns:
+        Tensor, with the designated shape and dtype, filled with random numbers from the uniform distribution on
+        the interval :math:`[low, high)`.
+
+    Raises:
+        TypeError: If `size` is not a tuple.
+        TypeError: If `low` or `high` is not integer.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore import ops
+        >>> print(ops.function.random_func.randint_ext(0, 5, (2, 3)).shape)
+        (2, 3)
+    """
+    if not generator:
+        generator = default_generator
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
+    return randint_(low, high, size, seed, offset, dtype)
+
+
+@_function_forbid_reuse
+def randint_like_ext(input, low, high, *, dtype=None):
+    r"""
+    Returns a new tensor filled with integer numbers from the uniform distribution over an interval :math:`[low, high)`
+    based on the given dtype and shape of the input tensor.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        input (Tensor): Input Tensor to specify the output shape and its default dtype.
+        low (int): the lower bound of the generated random number
+        high (int): the upper bound of the generated random number
+
+    Keyword Args:
+        dtype (:class:`mindspore.dtype`, optional): Designated tensor dtype. If None,
+            the same dtype of `input` will be applied. Default: ``None`` .
+
+    Returns:
+        Tensor, with the designated shape and dtype, filled with random numbers from the uniform distribution on
+        the interval :math:`[low, high)`.
+
+    Raises:
+        TypeError: If `low` or `high` is not integer.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import Tensor, ops
+        >>> a = Tensor([[2, 3, 4], [1, 2, 3]])
+        >>> low = 0
+        >>> high = 5
+        >>> print(ops.function.random_func.randint_like_ext(a, low, high, dtype=ms.int32).shape)
+        (2, 3)
+    """
+    seed, offset = default_generator._step(  # pylint: disable=protected-access
+        generator_step_)
+    return randint_like_(input, low, high, seed, offset, dtype)
+
+
 @_function_forbid_reuse
 def randn(*size, dtype=None, seed=None):
     r"""
@@ -1395,6 +1568,47 @@ def randint_like(input, low, high, seed=None, *, dtype=None):
     return cast_(output, dtype)
 
 
+def randperm_ext(n, *, generator=None, dtype=mstype.int64):
+    r"""
+    Generates random permutation of integers from 0 to n-1.
+
+    .. warning::
+        - This is an experimental API that is subject to change or deletion.
+
+
+    Args:
+        n (Union[Tensor, int]): size of the permutation. int or Tensor with shape: () or (1,) and
+            data type int64. The value of `n` must be greater than zero.
+        generator (:class:`mindspore.Generator`, optional): a pseudorandom number generator.
+            Default: ``None``, uses the default pseudorandom number generator.
+        dtype (mindspore.dtype, optional): The type of output. Default: mstype.int64.
+
+    Returns:
+        Tensor with shape (n,) and type `dtype`.
+
+    Raises:
+        TypeError: If `dtype` is not supported.
+        ValueError: If `n` is a negative or 0 element.
+        ValueError: If `n` is larger than the maximal data of the set dtype.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> from mindspore import ops
+        >>> from mindspore import dtype as mstype
+        >>> n = 4
+        >>> output = ops.randperm_ext(n, dtype=mstype.int64)
+        >>> print(output.shape)
+        (4,)
+    """
+    if not generator:
+        generator = default_generator
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
+    return randperm_ext_(n, seed, offset, dtype)
+
+
 @_function_forbid_reuse
 def poisson(shape, mean, seed=None):
     r"""
@@ -1675,10 +1889,10 @@ def multinomial_ext(input, num_samples, replacement=False, *, generator=None):
         >>> # [[0 0 0 0 0 0 0 0 1 0]
         >>> #  [1 1 1 1 1 0 1 1 1 1]]
     """
-
     if generator is None:
         generator = default_generator
-    seed, offset = generator._step(generator_step_)  # pylint: disable=protected-access
+    seed, offset = generator._step(  # pylint: disable=protected-access
+        generator_step_)
     return multinomial_ext_(input, num_samples, replacement, seed, offset)
 
 

mindspore/ops/operations/__init__.py

@@ -55,7 +55,7 @@ from .comm_ops import (AllGather, AllReduce, Reduce, NeighborExchange, NeighborE
                        _MirrorOperator, _MirrorMiniStepOperator, _MiniStepAllGather, ReduceOp, _VirtualDataset,
                        _VirtualOutput, _VirtualDiv, _GetTensorSlice, _VirtualAdd, _VirtualAssignAdd, _VirtualAccuGrad,
                        _HostAllGather, _HostReduceScatter, _MirrorMicroStepOperator, _MicroStepAllGather,
-                       _VirtualPipelineEnd, AlltoAllV, ReduceScatter)
+                       _VirtualPipelineEnd, AlltoAllV, ReduceScatter, _VirtualAssignKvCache)
 from .control_ops import GeSwitch, Merge
 from .custom_ops import (Custom)
 from .debug_ops import (ImageSummary, InsertGradientOf, HookBackward, ScalarSummary,

mindspore/ops/operations/array_ops.py

@@ -771,12 +771,14 @@ class Padding(Primitive):
 class UniqueWithPad(Primitive):
     """
     'ops.UniqueWithPad' is deprecated from version 2.4 and will be removed in a future version.
+    Please use the :func:`mindspore.ops.unique` combined with :func:`mindspore.ops.pad` to realize
+    the same function.
 
     Supported Platforms:
         Deprecated
     """
 
-    @deprecated("2.4", "ops.Unique and ops.PadV3", False)
+    @deprecated("2.4", "ops.unique and ops.pad", False)
     @prim_attr_register
     def __init__(self):
         """init UniqueWithPad"""

mindspore/ops/operations/comm_ops.py

@@ -1682,6 +1682,27 @@ class _VirtualAssignAdd(PrimitiveWithInfer):
 virtual_assign_add = _VirtualAssignAdd()
 
 
+class _VirtualAssignKvCache(PrimitiveWithInfer):
+    """
+    Auto parallel virtual operator. Do nothing in forward, do Assign kv cache in backward. It is only for
+    internal use of parallel modules and cannot be called by users.
+
+    """
+
+    @prim_attr_register
+    def __init__(self):
+        """Initialize _VirtualAssignAdd."""
+        self.add_prim_attr('order_enforce_skip', True)
+        self.add_prim_attr('side_effect_backprop_mem', True)
+
+    def infer_shape(self, x_shape, y_shape, kv_equal_shape):
+        return x_shape
+
+    def infer_dtype(self, x_dtype, y_dtype, kv_equal_dtype):
+        return x_dtype
+virtual_assign_kv_cache = _VirtualAssignKvCache()
+
+
 class _VirtualAccuGrad(PrimitiveWithInfer):
     """
     Auto parallel virtual operator. Do nothing in forward, return y in backward. It is only for

mindspore/ops/operations/manually_defined/ops_def.py

@@ -1171,17 +1171,15 @@ class Cast(Primitive):
         self.init_prim_io_names(inputs=['x', 'dst_type'], outputs=['output'])
 
     def check_elim(self, x, dtype):
-        if isinstance(x, (Tensor, numbers.Number, Parameter)):
-            if isinstance(x, Parameter):
-                data = x.data
-                if data.dtype == dtype:
-                    return (True, x)
-            if isinstance(x, Tensor) and x.dtype == dtype:
-                x = Tensor(x)
-                x.set_cast_dtype()
+        if isinstance(x, Parameter):
+            data = x.data
+            if data.dtype == dtype:
                 return (True, x)
-            if isinstance(x, numbers.Number):
-                return (True, Tensor(x, dtype=dtype))
+        if isinstance(x, Tensor) and x.dtype == dtype:
+            x.set_cast_dtype()
+            return (True, x)
+        if isinstance(x, numbers.Number):
+            return (True, Tensor(x, dtype=dtype))
         return (False, None)
 
     def __call__(self, input_x, dtype):

mindspore/parallel/_auto_parallel_context.py

@@ -76,6 +76,7 @@ class _PipelineConfig:
 class _PipelineScheduler:
     PIPELINE_1F1B = "1f1b"
     PIPELINE_GPIPE = "gpipe"
+    PIPELINE_SEQPIPE = "seqpipe"
 
 
 class _AutoParallelContext:
@@ -914,7 +915,8 @@ class _AutoParallelContext:
             pipeline_config[pp_interleave])
 
         Validator.check_string(pipeline_config[pp_scheduler], [_PipelineScheduler.PIPELINE_1F1B,
-                                                               _PipelineScheduler.PIPELINE_GPIPE])
+                                                               _PipelineScheduler.PIPELINE_GPIPE,
+                                                               _PipelineScheduler.PIPELINE_SEQPIPE])
         if not pipeline_config[pp_interleave] and pipeline_config[pp_scheduler] != _PipelineScheduler.PIPELINE_1F1B:
             raise ValueError(f"When pipeline_interleave is False, {pp_scheduler} is not supported")
 

mindspore/parallel/_cell_wrapper.py

@@ -126,6 +126,8 @@ def _restore_parallel_context(origin_parallel_mode, origin_dataset_strategy):
     if context.get_context("mode") == context.GRAPH_MODE:
         context.set_auto_parallel_context(parallel_mode=origin_parallel_mode)
         if origin_dataset_strategy != "data_parallel":
+            if origin_dataset_strategy is not None and isinstance(origin_dataset_strategy, list):
+                origin_dataset_strategy = tuple(tuple(ds_item) for ds_item in origin_dataset_strategy)
             context.set_auto_parallel_context(dataset_strategy=origin_dataset_strategy)