mindspore 2.6.0__cp311-cp311-win_amd64.whl → 2.7.0rc1__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
 
@@ -34,11 +32,10 @@ from mindspore._c_expression import pyboost_generator
 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):

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/comm_ops.py

@@ -300,7 +300,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.
@@ -957,25 +957,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 +1732,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 +1958,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 +1995,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 +2005,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 +2036,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)