mindspore 2.6.0rc1__cp311-cp311-win_amd64.whl → 2.7.0__cp311-cp311-win_amd64.whl

mindspore/ops/operations/_inner_ops.py

@@ -14,9 +14,7 @@
 # ============================================================================
 
 """Inner operators."""
-from types import FunctionType, MethodType
 from collections.abc import Iterable
-import os
 import weakref
 import numpy as np
 
@@ -27,18 +25,17 @@ from mindspore.ops.operations._scalar_ops import bit_or, bit_and
 from mindspore.ops import signature as sig
 from mindspore.ops.operations.math_ops import _infer_shape_reduce
 from mindspore.ops.primitive import PrimitiveWithCheck, PrimitiveWithInfer, prim_attr_register, Primitive, \
-    _run_op, _check_contains_variable
+    _check_contains_variable
 from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._c_expression import typing, HookType
-from mindspore._c_expression import pyboost_generator
+from mindspore._c_expression import pyboost_generator, pyboost_cell_backward_hook
 from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.common.parameter import Parameter
-from mindspore.common._stub_tensor import _convert_stub
 from mindspore.communication.management import GlobalComm, get_rank, _get_group, get_group_size
 from mindspore.common.api import _pynative_executor
-from ..auto_generate import TensorCopySlices, SiLU, Cummin, TopKRouter, ExtractImagePatches, DecoderKVCache, \
-    PromptKVCache, ApplyCamePart1, ApplyCamePart2, ApplyCamePart3, ApplyCamePart4
+from ..auto_generate import TensorCopySlices, SiLU, Cummin, TopKRouter, TopPRouter, ExtractImagePatches, \
+    ApplyCamePart1, ApplyCamePart2, ApplyCamePart3, ApplyCamePart4
 
 # Bit operation
 bit_and = bit_and()
@@ -78,7 +75,7 @@ class Generator(Primitive):
     def __call__(self, cmd, inputs):
         if cmd == 0:  # step cmd
             return inputs[0], inputs[1]
-        return _convert_stub(pyboost_generator(self, [cmd, inputs]))
+        return pyboost_generator(self, [cmd, inputs])
 
 
 class Quant(PrimitiveWithInfer):
@@ -1568,10 +1565,11 @@ class CellBackwardHook(PrimitiveWithInfer):
         self.add_prim_attr("cell_id", cell_id)
         self.grad_output = None
 
-    def __call__(self, *args):
-        # If args is empty, just return.
-        if not args:
-            return args
+    def __call__(self, args):
+        is_tuple = True
+        if not isinstance(args, tuple):
+            args = (args,)
+            is_tuple = False
 
         # Collect the indices and values of arguments that are instances of Tensor
         tensors_idx = []
@@ -1581,20 +1579,18 @@ class CellBackwardHook(PrimitiveWithInfer):
                 tensors_idx.append(i)
                 tensors.append(arg)
 
-        # If there are no Tensor arguments, return the single argument or the original tuple
+        # If there are no Tensor arguments, return original args
         if not tensors:
-            return args[0] if len(args) == 1 else args
+            return args if is_tuple else args[0]
 
-        new_tensors = _run_op(self, self.name, tensors)
-        if not isinstance(new_tensors, tuple):
-            new_tensors = (new_tensors,)
+        new_tensors = pyboost_cell_backward_hook(self, (tensors,))
 
         # Replace the original Tensor arguments with the processed ones
         arg_list = list(args)
         for idx, val in zip(tensors_idx, new_tensors):
             arg_list[idx] = val
 
-        return arg_list[0] if len(arg_list) == 1 else tuple(arg_list)
+        return tuple(arg_list) if is_tuple else arg_list[0]
 
     def infer_shape(self, *inputs_shape):
         if len(inputs_shape) == 1:

mindspore/ops/operations/_sequence_ops.py

@@ -515,7 +515,7 @@ class TensorToScalar(PrimitiveWithCheck):
         """infer_value TensorToScalar"""
         if isinstance(x, (Tensor_, Tensor)):
             if not x.shape or x.shape == (1,) or x.shape == [1,]:
-                return x.asnumpy().item()
+                return x.item()
         return None
 
 

mindspore/ops/operations/array_ops.py

@@ -36,13 +36,13 @@ from mindspore._c_expression import TensorPy as Tensor_
 from mindspore._c_expression import CSRTensor as CSRTensor_
 from mindspore._c_expression import COOTensor as COOTensor_
 from ..auto_generate import (
-    ExpandDims, Reshape, TensorShape, Transpose, Gather, OnesLike, ZerosLike,
+    ExpandDims, Reshape, TensorShape, Transpose, TransposeView, Gather, OnesLike, ZerosLike,
     Argmax, ArgMaxExt, ReverseV2, Diag, Eye, ScatterNd,
     ResizeNearestNeighborV2, GatherNd, GatherD, Range, MaskedFill, RightShift,
     NonZero, ResizeNearestNeighbor, Identity, Split, CumSum, CumProd,
     MaskedSelect, Cummax, Cummin, Argmin, Concat, UnsortedSegmentSum, UniqueConsecutive,
     ScalarToTensor, Triu, BroadcastTo, StridedSlice, Select, TopkExt,
-    SearchSorted, Meshgrid, Squeeze, Slice)
+    SearchSorted, Meshgrid, Squeeze, Slice, TransposeExtView, MaskedScatter)
 from .manually_defined import Rank, Shape, Tile, Cast, Ones, Zeros, TypeAs
 from ..auto_generate import ArgMaxWithValue, ArgMinWithValue
 from ..auto_generate import TensorScatterElements as TensorScatterElementsExt
@@ -1048,11 +1048,11 @@ class Fill(PrimitiveWithCheck):
         self.init_prim_io_names(inputs=['type', 'shape', 'value'], outputs=['y'])
 
     def __call__(self, dtype, dims, x):
-        if dtype not in mstype.all_types and dtype not in [mstype.uint16, mstype.uint32, mstype.uint64]:
+        if dtype not in mstype.all_types:
             raise TypeError(
                 f"For \'{self.name}\', the supported data type is ['bool', 'int8', 'int16', 'int32', 'int64', 'uint8', "
                 "'uint16', 'uint32', 'uint64','float16', 'float32', 'float64'], but got an invalid dtype!.")
-        x_nptype = mstype.dtype_to_nptype(dtype)
+        x_nptype = mstype._dtype_to_nptype(dtype)  # pylint:disable=protected-access
         if not isinstance(dims, Tensor) and not isinstance(dims, tuple):
             raise TypeError(f"For \'{self.name}\', input[1] must be tensor.")
         if not isinstance(x, Tensor) and not isinstance(x, float) and not isinstance(x, int):
@@ -1065,7 +1065,7 @@ class Fill(PrimitiveWithCheck):
         return Tensor(ret, dtype=dtype)
 
     def infer_value(self, dtype, dims, x):
-        x_nptype = mstype.dtype_to_nptype(dtype)
+        x_nptype = mstype._dtype_to_nptype(dtype)  # pylint:disable=protected-access
         if dims is not None and None not in dims and x is not None:
             if isinstance(dims, Tensor):
                 dims = dims.asnumpy()
@@ -3974,52 +3974,6 @@ class RangeV2(Primitive):
         validator.check_positive_int(maxlen, "maxlen", self.name)
 
 
-class MaskedScatter(Primitive):
-    """
-    Updates the value in the input with value in `updates` according to the `mask`.
-
-    .. warning::
-        This is an experimental API that is subject to change or deletion.
-
-    Inputs:
-        - **x** (Tensor): The input Tensor to be updated.
-        - **mask** (Tensor[bool]): The mask Tensor indicating which elements should be modified or replaced.
-          The shapes of `mask` and `x` must be the same or broadcastable.
-        - **updates** (Tensor): The values to scatter into the target tensor `x`. It has the same data type as `x`. The
-          number of elements must be greater than or equal to the number of True's in `mask`.
-
-    Outputs:
-        Tensor, with the same type and shape as `x`.
-
-    Raises:
-        TypeError: If `x`, `mask` or `updates` is not a Tensor.
-        TypeError: If data type of `x` is not be supported.
-        TypeError: If dtype of `mask` is not bool.
-        TypeError: If the dim of `x` less than the dim of `mask`.
-        ValueError: If `mask` can not be broadcastable to `x`.
-        ValueError: If the number of elements in `updates` is less than number of True's in `mask`.
-
-    Supported Platforms:
-        ``Ascend`` ``CPU``
-
-    Examples:
-        >>> import mindspore
-        >>> import numpy as np
-        >>> from mindspore import Tensor, ops
-        >>> input_x = Tensor(np.array([1., 2., 3., 4.]), mindspore.float32)
-        >>> mask = Tensor(np.array([True, True, False, True]), mindspore.bool_)
-        >>> updates = Tensor(np.array([5., 6., 7.]), mindspore.float32)
-        >>> output = ops.MaskedScatter()(input_x, mask, updates)
-        >>> print(output)
-        [5. 6. 3. 7.]
-    """
-
-    @prim_attr_register
-    def __init__(self):
-        """Initialize MaskedScatter"""
-        self.init_prim_io_names(inputs=['x', 'mask', 'updates'], outputs=['y'])
-
-
 class _TensorScatterOp(PrimitiveWithInfer):
     """
     Defines TensorScatter Base Operators

mindspore/ops/operations/comm_ops.py

@@ -18,10 +18,9 @@
 from __future__ import absolute_import
 from __future__ import division
 
-import os
 from mindspore.common import Tensor
 from mindspore import _checkparam as validator
-from mindspore.communication.management import get_rank, get_group_size, GlobalComm, _get_group, _host_distribute
+from mindspore.communication.management import get_rank, get_group_size, GlobalComm, _get_group
 from mindspore.common import dtype as mstype
 from mindspore.ops.primitive import PrimitiveWithInfer, PrimitiveWithCheck, Primitive, prim_attr_register
 from mindspore.common.api import context
@@ -98,17 +97,6 @@ def check_collective_target_dtype(data_name, data_dtype, prim_name):
     validator.check_tensor_dtype_valid(data_name, data_dtype, valid_dtype, prim_name)
 
 
-def check_hcom_group_valid(group, prim_name=None):
-    """Check if hcom group is valid."""
-    msg_prefix = f"For '{prim_name}', the" if prim_name else "The"
-    sim_level = os.getenv("MS_SIMULATION_LEVEL")
-    no_sim = (sim_level is None or sim_level.strip() == '')
-    if no_sim and (not _host_distribute()) and context.get_context("mode") == context.PYNATIVE_MODE and \
-            group != GlobalComm.WORLD_COMM_GROUP:
-        raise RuntimeError(f"{msg_prefix} 'group' only support 'hccl_world_group' in pynative mode, but got "
-                           f"'group': {group}. Please start by using mpi-run.")
-
-
 class AllReduce(Primitive):
     """
     Reduces tensors across all devices in such a way that all devices will get the same final result,
@@ -187,7 +175,6 @@ class AllReduce(Primitive):
         if not isinstance(self.group, str):
             raise TypeError(f"For '{self.name}', the 'group' must be str, "
                             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', self.group)
         self.add_prim_attr('fusion', 0)
@@ -300,7 +287,7 @@ class AllGather(PrimitiveWithInfer):
 
     Outputs:
         Tensor. If the number of devices in the group is N,
-        then the shape of output is :math:`(N, x_1, x_2, ..., x_R)`.
+        then the shape of output is :math:`(N*x_1, x_2, ..., x_R)`.
 
     Raises:
         TypeError: If `group` is not a str.
@@ -720,7 +707,6 @@ class Broadcast(PrimitiveWithInfer):
         """Initialize Broadcast."""
         validator.check_value_type('root_rank', root_rank, (int,), self.name)
         validator.check_value_type('group', _get_group(group), (str,), self.name)
-        check_hcom_group_valid(group, prim_name=self.name)
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('no_eliminate', True)
 
@@ -957,25 +943,8 @@ class AlltoAll(PrimitiveWithInfer):
         self.concat_dim = concat_dim
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('no_eliminate', True)
-
-    def infer_shape(self, x_shape):
-        rank_size = get_group_size(_get_group(self.group))
-        if self.split_count != rank_size:
-            raise ValueError(f"For '{self.name}', the 'split_count' must be equal to 'rank_size', "
-                             f"but got 'split_count': {self.split_count}, 'rank_size': {rank_size}.")
-        if x_shape[self.split_dim] >= 0 and x_shape[self.split_dim] % self.split_count != 0:
-            raise ValueError(f"For '{self.name}', the 'x_shape[self.split_dim]' must be divisible by 'split_count', "
-                             f"but got 'x_shape[self.split_dim]' {x_shape[self.split_dim]}, "
-                             f"'split_count' {self.split_count}.")
-        if x_shape[self.concat_dim] >= 0:
-            x_shape[self.concat_dim] = x_shape[self.concat_dim] * self.split_count
-        if x_shape[self.split_dim] >= 0:
-            x_shape[self.split_dim] = int(x_shape[self.split_dim] / self.split_count)
-        return x_shape
-
-    def infer_dtype(self, x_dtype):
-        check_collective_target_dtype('x', x_dtype, self.name)
-        return x_dtype
+        self.rank_size = get_group_size(_get_group(group))
+        self.add_prim_attr('rank_size', self.rank_size)
 
 
 class NeighborExchangeV2(Primitive):
@@ -1749,10 +1718,10 @@ class _VirtualAssignKvCache(PrimitiveWithInfer):
         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):
+    def infer_shape(self, x_shape, y_shape, kv_equal_shape=None):
         return x_shape
 
-    def infer_dtype(self, x_dtype, y_dtype, kv_equal_dtype):
+    def infer_dtype(self, x_dtype, y_dtype, kv_equal_dtype=None):
         return x_dtype
 virtual_assign_kv_cache = _VirtualAssignKvCache()
 
@@ -1975,8 +1944,8 @@ class AlltoAllV(PrimitiveWithInfer):
     AllToAllV which support uneven scatter and gather compared with AllToAll.
 
     Note:
-        - Only support flatten tensor as input. input tensor should be flattened and
-          concatenated before call this primitive.
+        Only support flatten tensor as input. input tensor should be flattened and
+        concatenated before call this primitive.
 
     Args:
         group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP``, which
@@ -2012,9 +1981,9 @@ class AlltoAllV(PrimitiveWithInfer):
 
             This example should be run with 2 devices.
 
-        >>> from mindspore import ops
         >>> import mindspore.nn as nn
         >>> from mindspore.communication import init, get_rank
+        >>> from mindspore.ops import AlltoAllV
         >>> from mindspore import Tensor
         >>>
         >>> init()
@@ -2022,7 +1991,7 @@ class AlltoAllV(PrimitiveWithInfer):
         >>> class Net(nn.Cell):
         ...     def __init__(self):
         ...         super(Net, self).__init__()
-        ...         self.all_to_all = ops.AlltoAllV()
+        ...         self.all_to_all = AlltoAllV()
         ...
         ...     def construct(self, x, send_numel_list, recv_numel_list):
         ...         return self.all_to_all(x, send_numel_list, recv_numel_list)
@@ -2053,3 +2022,179 @@ class AlltoAllV(PrimitiveWithInfer):
         self.add_prim_attr('group', self.group)
         validator.check_value_type("block_size", block_size, [int], self.name)
         self.add_prim_attr('block_size', self.block_size)
+
+
+class AllGatherV(PrimitiveWithInfer):
+    """
+    Gathers uneven tensors from the specified communication group and returns the tensor which is all gathered.
+
+    Note:
+        Only support flatten tensor as input. input tensor should be flattened and
+        concatenated before call this primitive.
+
+    Args:
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
+            means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
+
+    Inputs:
+        - **input_x** (Tensor) - One-dimensional tensor to be gathered, with the shape :math:`(x_1)`.
+        - **output_split_sizes** (Union[tuple[int], list[int], Tensor]) - One-dimensional tensor, a list of the
+          amount of data gathered by all ranks. The basic unit is the data type of Tensor.
+
+    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.
+
+    Raises:
+        RuntimeError: Device target is invalid, backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    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.
+
+        >>> import mindspore as ms
+        >>> from mindspore.ops import AllGatherV
+        >>> import mindspore.nn as nn
+        >>> from mindspore.communication import init, get_rank
+        >>> from mindspore import Tensor
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> init()
+        >>> class Net(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net, self).__init__()
+        ...         self.allgatherv = AllGatherV()
+        ...
+        ...     def construct(self, x, output_split_sizes):
+        ...         return self.allgatherv(x, output_split_sizes)
+        ...
+        >>> rank = get_rank()
+        >>> data = [i for i in range(rank + 3)]
+        >>> input_x = Tensor(data)
+        >>> output_split_sizes = [3, 4]
+        >>> net = Net()
+        >>> output = net(input_x, output_split_sizes)
+        >>> print(output)
+        [0 1 2 0 1 2 3]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - AllGatherV
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#allgatherv>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, group=GlobalComm.WORLD_COMM_GROUP):
+        """Initialize AllGatherV."""
+        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', self.group)
+        self.add_prim_attr('mean_flag', False)
+        self.rank_id = get_rank(_get_group(self.group))
+        self.add_prim_attr('rank_id', self.rank_id)
+
+
+class ReduceScatterV(PrimitiveWithInfer):
+    r"""
+    Reduces and scatters uneven tensors from the specified communication group
+    and returns the tensor which is reduced and scattered.
+
+    Note:
+        Only support flatten tensor as input. The input tensor should be flattened and
+        concatenated before call this primitive.
+
+    Args:
+        op (str, optional): Specifies an operation used for element-wise reductions,
+            like SUM, MIN and MAX, currently PROD is not supported. Default: ``ReduceOp.SUM`` .
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` , which
+            means ``"hccl_world_group"`` in Ascend, and ``"nccl_world_group"`` in GPU.
+
+    Inputs:
+        - **input_x** (Tensor) - One-dimensional tensor to be distributed, with the shape :math:`(x_1)`.
+        - **input_split_sizes** (Union[tuple[int], list[int], Tensor]) - One-dimensional tensor, a list of
+          received data volumes for all ranks. The basic unit is the data type of Tensor. The value is not
+          verified, and the user guarantees its correctness.
+
+    Outputs:
+        Tensor. Reduces and scatters tensor from remote ranks.
+        If the result is empty, it will return a Tensor with shape `()`, and value has no actual meaning.
+
+    Raises:
+        RuntimeError: Device target is invalid, backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU``
+
+    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.
+
+        >>> import mindspore as ms
+        >>> from mindspore import Tensor
+        >>> from mindspore.communication import init, get_rank
+        >>> from mindspore.ops import ReduceOp
+        >>> import mindspore.nn as nn
+        >>> from mindspore.ops.operations.comm_ops import ReduceScatterV
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
+        >>> init()
+        >>> class Net(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net, self).__init__()
+        ...         self.reducescatterv = ReduceScatterV(ReduceOp.SUM)
+        ...
+        ...     def construct(self, x, input_split_sizes):
+        ...         return self.reducescatterv(x, input_split_sizes)
+        ...
+        >>> rank = get_rank()
+        >>> input_x = Tensor([0, 1, 2.0])
+        >>> input_split_sizes = [2, 1]
+        >>> net = Net()
+        >>> output = net(input_x, input_split_sizes)
+        >>> print(output)
+        rank 0:
+        [0. 2.]
+        rank 1:
+        [4.]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - ReduceScatterV
+          <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#reducescatterv>`_
+
+    """
+
+    @prim_attr_register
+    def __init__(self, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
+        """Initialize ReduceScatterV."""
+        validator.check_value_type('op', op, (type(ReduceOp.SUM),), 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(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)