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

mindspore/mint/distributed/distributed.py

@@ -58,9 +58,11 @@ from mindspore.ops.auto_generate.gen_ops_prim import (
     dist_comm_isend_op,
     dist_comm_all_to_all_v_op,
     dist_comm_reduce_scatter_tensor_op,
+    dist_comm_reduce_scatter_tensor_uneven_op,
     dist_comm_all_to_all_v_single_op,
     dist_comm_broadcast_op,
     dist_comm_all_gather_into_tensor_op,
+    dist_comm_all_gather_into_tensor_uneven_op,
     dist_comm_irecv_op,
     dist_comm_scatter_tensor_op,
     dist_comm_gather_into_tensor_op,
@@ -70,6 +72,7 @@ from mindspore.ops.auto_generate.gen_ops_prim import (
     dist_comm_barrier_op,
     dist_comm_batch_isend_irecv_op,
 )
+from mindspore._c_expression import TCPStoreClient, GroupOptions
 
 _pickler = pickle.Pickler
 _unpickler = pickle.Unpickler
@@ -137,6 +140,193 @@ def _tensor_to_object(tensor, tensor_size):
     return restricted_loads(buf)
 
 
+class TCPStore:
+    """
+    A TCP-based distributed key-value store implementation.
+
+    Note:
+        - The function is implemented by CPU and does not involve any hardware operations related to Ascend.
+        - Currently, all parameters provided by the TCPStore class constructor are not supported.
+          The master node and port number are uniformly specified by the MindSpore framework.
+          The following parameters are provided, currently not supported and settings are invalid.
+        - The current TcpStore function is limited and only supports scenarios where the key is
+          less than 4k and the value is less than 1G. Complex scenarios are to be supported.
+        - The timeout interval for message sending and receiving in the TcpStore function is controlled by
+          the `MS_RECEIVE_MSG_TIMEOUT` environment variable, in seconds, with a default value of ``15``.
+          If a timeout occurs, the user needs to increase the configuration value.
+
+    Args:
+        host_name (str, invalid, optional): The hostname or IP Address the server store should run on.
+            Default is ``None``.
+        port (int, invalid, optional): The port on which the server store should listen for incoming requests.
+            Default is ``None``.
+        world_size (int, invalid, optional): The total number of store users (number of clients + 1 for the server).
+            Default is ``None`` (``None`` indicates a non-fixed number of store users).
+        is_master (bool, invalid, optional): True when initializing the server store and False for client stores.
+            Default is ``False``.
+        timeout (timedelta, invalid, optional): Timeout used by the store during initialization, Unit: seconds.
+            Default is ``300``.
+        wait_for_workers (bool, invalid, optional): Whether to wait for all the workers to connect with the server
+            store. This is only applicable when `world_size` is a fixed value. Default is ``True``.
+        multi_tenant (bool, invalid, optional): If ``True``, all ``TCPStore`` instances in the current process with
+            the same host/port will use the same underlying ``TCPServer``. Default is ``False``.
+        master_listen_fd (int, invalid, optional): If specified, the underlying ``TCPServer`` will listen on this file
+            descriptor, which must be a socket already bound to ``port``. Useful to avoid port assignment races
+            in some scenarios. Default is ``None`` (meaning the server creates a new socket and attempts to bind it
+            to ``port``).
+        use_libuv (bool, invalid, optional): If True, use libuv for ``TCPServer`` backend. Default is ``True``.
+
+    Returns:
+        TCPStore Object.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend 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.
+
+        >>> from mindspore.mint.distributed import TCPStore
+        >>> store = TCPStore()
+    """
+
+    def __init__(self, host_name=None, port=None, world_size=None, is_master=False, timeout=300,
+                 wait_for_workers=True, multi_tenant=False, master_listen_fd=None, use_libuv=True):
+        self.instance = TCPStoreClient.get_instance()
+
+
+    def set(self, key, value):
+        """
+        Inserts the key-value pair into the store based on the supplied `key` and
+        `value`. If `key` already exists in the store, it will overwrite the old
+        value with the new supplied `value`.
+
+        Args:
+            key (str): The key to be added to the store.
+            value (Union[bytes, str]): The value associated with `key` to be added to the store.
+
+        Raises:
+            TypeError: If `key` is not string.
+            TypeError: If `value` is not string or bytes.
+
+        Supported Platforms:
+            ``Ascend``
+
+        Examples:
+            .. note::
+                Before running the following examples, you need to configure the communication environment variables.
+
+                For Ascend 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.
+
+            >>> from mindspore.mint.distributed import TCPStore
+            >>> store = TCPStore()
+            >>> store.set("first_key", "first_value")
+        """
+        if not isinstance(key, str):
+            raise TypeError(
+                "For 'TCPStore.set', the argument 'key' must be type of string, "
+                "but got 'key' type : {}.".format(type(key))
+            )
+        if not isinstance(value, (str, bytes)):
+            raise TypeError(
+                "For 'TCPStore.set', the argument 'value' must be type of string or bytes, "
+                "but got 'value' type : {}.".format(type(value))
+            )
+        return self.instance.set(key, value)
+
+
+    def get(self, key):
+        """
+        Retrieves the value associated with the given `key` in the store. If `key` is not
+        present in the store, the function will return "".
+
+        Args:
+            key (str): The function will return the value associated with this key.
+
+        Returns:
+            bytes, Value associated with `key` if `key` is in the store.
+
+        Raises:
+            TypeError: If `key` is not string.
+
+        Supported Platforms:
+            ``Ascend``
+
+        Examples:
+            .. note::
+                Before running the following examples, you need to configure the communication environment variables.
+
+                For Ascend 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.
+
+            >>> from mindspore.mint.distributed import TCPStore
+            >>> store = TCPStore()
+            >>> store.set("first_key", "first_value")
+            >>> data = store.get("first_key")
+            >>> print(data)
+        """
+        if not isinstance(key, str):
+            raise TypeError(
+                "For 'TCPStore.get', the argument 'key' must be type of string, "
+                "but got 'key' type : {}.".format(type(key))
+            )
+        byte_data = self.instance.get(key)
+        return byte_data
+
+
+    def delete_key(self, key):
+        """
+        Deletes the key-value pair associated with `key` from the store.
+
+        Args:
+            key (str): The key to be deleted from the store.
+
+        Returns:
+            bool, ``True`` if `key` was deleted, otherwise ``False``.
+
+        Raises:
+            TypeError: If `key` is not string.
+
+        Supported Platforms:
+            ``CPU``
+
+        Examples:
+            .. note::
+                Before running the following examples, you need to configure the communication environment variables.
+
+                For Ascend 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.
+
+            >>> from mindspore.mint.distributed import TCPStore
+            >>> store = TCPStore()
+            >>> store.set("first_key", "first_value")
+            >>> # This should return true
+            >>> store.delete_key("first_key")
+        """
+        if not isinstance(key, str):
+            raise TypeError(
+                "For 'TCPStore.delete_key', the argument 'key' must be type of string, "
+                "but got 'key' type : {}.".format(type(key))
+            )
+        return self.instance.delete_key(key)
+
+
 def is_available():
     """
     Checks if distributed module is available.
@@ -470,7 +660,19 @@ def new_group(ranks=None,
             when backend is ``"hccl"`` will use Huawei Collective Communication Library(HCCL).
             when  backend is ``"mccl"`` will use MindSpore Collective Communication Library(MCCL).
             If ``None``, which means ``"hccl"`` in Ascend. Default is ``None``.
-        pg_options (str, invalid): Currently it is a reserved parameter.
+        pg_options (GroupOptions, optional): Additional communication group configuration parameters.
+            The backend will automatically select supported parameters and apply them during group
+            initialization. i.e. for the ``HCCL`` backend, ``hccl_config`` can be specified so that
+            group initialization configurations can be applied. Default is ``None``.
+
+            `GroupOptions` is defined as a class that can be instantiated as a python object.
+
+            .. code-block::
+
+                GroupOptions {
+                    hccl_config(dict)
+                }
+
         use_local_synchronization (bool, invalid): Currently it is a reserved parameter.
         group_desc (str, invalid): Currently it is a reserved parameter.
 
@@ -511,8 +713,11 @@ def new_group(ranks=None,
     if not isinstance(backend, str) or backend not in ("hccl", "mccl"):
         raise TypeError(f"the input backend must be hccl or mccl, but got {backend}")
     group = backend + "_" + str(len(ranks)) + "_" + hashlib.sha1(bytes("_".join(map(str, ranks)), "utf-8")).hexdigest()
+    if pg_options is not None:
+        if not isinstance(pg_options, GroupOptions):
+            raise TypeError("pg_options must be type GroupOptions, but got {}".format(type(pg_options)))
     try:
-        create_group(group, ranks)
+        create_group(group, ranks, pg_options)
     except RuntimeError as e:
         logger.warning(e)
         group = ""
@@ -786,6 +991,22 @@ def _check_all_tensor_same_dtype_and_shape(*tensor_lists):
                     )
 
 
+@_primexpr
+def _check_output_shape(output, expected_shape, op_name):
+    if output.shape != expected_shape:
+        raise TypeError(
+            f"For {op_name}, the output shape should be {expected_shape}, "
+            f"but got {output.shape}.")
+
+
+@_primexpr
+def _check_output_dtype(output, expected_dtype, op_name):
+    if output.dtype != expected_dtype:
+        raise TypeError(
+            f"For {op_name}, the output dtype should be {expected_dtype}, "
+            f"but got {output.dtype}.")
+
+
 def all_reduce(tensor, op=ReduceOp.SUM, group=None, async_op=False):
     """
     Reduce tensors across all devices in such a way that all deviceswill get the same final result,
@@ -813,7 +1034,7 @@ def all_reduce(tensor, op=ReduceOp.SUM, group=None, async_op=False):
         RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
-        ``Ascend``
+        ``Ascend`` ``CPU``
 
     Examples:
         .. note::
@@ -950,6 +1171,91 @@ def all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=Fal
     return handle
 
 
+def all_gather_into_tensor_uneven(output, input, output_split_sizes=None, group=None, async_op=False):
+    r"""
+    Gathers and concatenates tensors across devices with uneven first dimensions.
+
+    Note:
+        - Input tensors must have identical shapes except for the first dimension.
+        - Output tensor's first dimension should equal to the sum of all devices' input first dimensions.
+
+    Args:
+        output (Tensor): Concatenated output tensor with shape :math:`(\sum_{i=0}^{N-1} x_{i1}, x_2, ..., x_R)`,
+            where N is the number of devices in the group.
+        input (Tensor): Local input tensor with shape :math:`(x_{k1}, x_2, ..., x_R)`, where k is current device's rank.
+        output_split_sizes (list[int], optional): Specifies first dimension sizes from each device.
+            Must match actual input dimensions when provided.
+            If ``None``, assumes equal split sizes across devices. Default: ``None``.
+        group (str, optional): The communication group to work on. If ``None``,
+            which means ``"hccl_world_group"`` in Ascend. Default: ``None``.
+        async_op (bool, optional): Whether this operator should be an async operator. Default: ``False``.
+
+    Returns:
+        CommHandle, CommHandle is an async work handle, if `async_op` is set to True.
+        CommHandle will be None, when `async_op` is False.
+
+    Raises:
+        ValueError: If the shape of `input` does not match the constraints of `output_split_sizes`.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend 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 numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore import ops
+        >>> from mindspore.mint.distributed import init_process_group, get_rank
+        >>> from mindspore.mint.distributed import all_gather_into_tensor_uneven
+        >>> from mindspore import Tensor
+        >>>
+        >>> ms.set_device(device_target="Ascend")
+        >>> init_process_group()
+        >>> if get_rank() == 0:
+        >>>     input_tensor = Tensor(np.ones([3, 4]).astype(np.float32))
+        >>> else:
+        >>>     input_tensor = Tensor(np.ones([2, 4]).astype(np.float32))
+        >>> out_tensor = Tensor(np.zeros([5, 4]).astype(np.float32))
+        >>> output_split_sizes = [3, 2]
+        >>> output = all_gather_into_tensor_uneven(out_tensor, input_tensor, output_split_sizes)
+        >>> print(out_tensor)
+        [[1. 1. 1. 1.]
+         [1. 1. 1. 1.]
+         [1. 1. 1. 1.]
+         [1. 1. 1. 1.]
+         [1. 1. 1. 1.]]
+    """
+    if group is None:
+        group = GlobalComm.WORLD_COMM_GROUP
+    if not isinstance(group, str):
+        raise TypeError(
+            "The argument 'group' must be type of string, "
+            "but got 'group' type : {}.".format(type(group))
+        )
+    if not isinstance(async_op, bool):
+        raise TypeError(
+            f"The argument 'async_op' must be a bool, but got {type(async_op)}."
+        )
+    group_size = get_cache_group_size(group)
+    output_split_sizes = [] if output_split_sizes is None else output_split_sizes
+    result = dist_comm_all_gather_into_tensor_uneven_op(
+        output, input, output_split_sizes, group_size, group
+    )
+    _, handle = _deal_comm_outputs(result, async_op)
+    return handle
+
+
 def reduce_scatter_tensor(output, input, op=ReduceOp.SUM, group=None, async_op=False):
     r"""
     Reduces and scatters tensors from the specified communication group and
@@ -1040,6 +1346,101 @@ def reduce_scatter_tensor(output, input, op=ReduceOp.SUM, group=None, async_op=F
     return handle
 
 
+def reduce_scatter_tensor_uneven(output, input, input_split_sizes=None, op=ReduceOp.SUM, group=None, async_op=False):
+    r"""
+    Reduce tensors from the specified communication group and scatter to the output tensor
+    according to `input_split_sizes`.
+
+    Note:
+        - The input tensor must have identical shape and format across all processes.
+        - The first dimension of input tensor should equal to the sum of `input_split_sizes`.
+
+    Args:
+        output(Tensor): the output tensor has the same dtype as `input` with a shape of
+            :math:`(input_split_sizes[rank], *)`, where rank is the local rank id of the device.
+        input(Tensor): The input tensor to be reduced and scattered, Expected shape :math:`(N, *)`, where `*`
+            means any number of additional dimensions. N must equal the sum of `input_split_sizes` across ranks.
+        input_split_sizes (list[int], optional): List specifying how to split the first dimension of input tensor.
+            If ``None``, splits evenly according to group size. Default: ``None``.
+        op (str, optional): Specifies an operation used for element-wise reductions,
+            One of ReduceOp: 'SUM', 'MIN', 'MAX'. Default: ``ReduceOp.SUM``.
+        group (str, optional): The communication group to work on. If ``None``, which means ``"hccl_world_group"`` in
+            Ascend. Default: ``None``.
+        async_op (bool, optional): Whether this operator should be an async operator. Default: ``False``.
+
+    Returns:
+        CommHandle, CommHandle is an async work handle, if `async_op` is set to True.
+        CommHandle will be None, when `async_op` is False.
+
+    Raises:
+        ValueError: If the shape of `output` does not match the constraints of `input_split_sizes`.
+        RuntimeError: If device target is invalid, or backend is invalid, or distributed initialization fails.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        .. note::
+            Before running the following examples, you need to configure the communication environment variables.
+
+            For Ascend 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.mint.distributed import init_process_group, get_rank
+        >>> from mindspore.mint.distributed import reduce_scatter_tensor_uneven
+        >>> import numpy as np
+        >>>
+        >>> ms.set_device(device_target="Ascend")
+        >>> init_process_group()
+        >>> input_tensor = Tensor(np.ones([5, 8]).astype(np.float32))
+        >>> if get_rank() == 0:
+        >>>     output_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> else:
+        >>>     output_tensor = Tensor(np.ones([3, 8]).astype(np.float32))
+        >>> input_split_sizes = [2, 3]
+        >>> output = reduce_scatter_tensor_uneven(output_tensor, input_tensor, input_split_sizes)
+        >>> print(output_tensor)
+        rank 0:
+        [[2. 2. 2. 2. 2. 2. 2. 2.]
+         [2. 2. 2. 2. 2. 2. 2. 2.]]
+        rank 1:
+        [[2. 2. 2. 2. 2. 2. 2. 2.]
+         [2. 2. 2. 2. 2. 2. 2. 2.]
+         [2. 2. 2. 2. 2. 2. 2. 2.]]
+    """
+    if not isinstance(op, str):
+        raise TypeError("For reduce_scatter_tensor_uneven, the input op type must be str")
+    if op not in ("sum", "min", "max"):
+        raise TypeError(
+            "For reduce_scatter_tensor_uneven, the input op value must be one of sum, prod, min, max"
+        )
+    if group is None:
+        group = GlobalComm.WORLD_COMM_GROUP
+    if not isinstance(group, str):
+        raise TypeError(
+            "The argument 'group' must be type of string, "
+            "but got 'group' type : {}.".format(type(group))
+        )
+    if not isinstance(async_op, bool):
+        raise TypeError(
+            f"The argument 'async_op' must be a bool, but got {type(async_op)}."
+        )
+    input_split_sizes = [] if input_split_sizes is None else input_split_sizes
+    rank_size = get_cache_group_size(group)
+    result = dist_comm_reduce_scatter_tensor_uneven_op(
+        output, input, input_split_sizes, rank_size, op, group
+    )
+    _, handle = _deal_comm_outputs(result, async_op)
+    return handle
+
+
 def reduce(tensor, dst, op=ReduceOp.SUM, group=None, async_op=False):
     """
     Reduces tensors across the processes in the specified communication group, sends the result
@@ -1589,7 +1990,7 @@ def barrier(group=None, async_op=False, device_ids=None):
         RuntimeError: If backend is invalid, or distributed initialization fails.
 
     Supported Platforms:
-        ``Ascend``
+        ``Ascend`` ``CPU``
 
     Examples:
         .. note::
@@ -1647,7 +2048,7 @@ def send(tensor, dst=0, group=None, tag=0):
         ValueError: If the `dst` process rank id is same as the current process.
 
     Supported Platforms:
-        ``Ascend``
+        ``Ascend`` ``CPU``
 
     Examples:
         .. note::
@@ -1725,7 +2126,7 @@ def recv(tensor, src=0, group=None, tag=0):
         ValueError: If the rank ID of the process is greater than the rank size of the communication group.
 
     Supported Platforms:
-        ``Ascend``
+        ``Ascend`` ``CPU``
 
     Examples:
         .. note::
@@ -2018,10 +2419,10 @@ def all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False
     recv_shape_list = []
 
     for tensor in input_tensor_list:
-        send_numel_list.append(tensor.size)
+        send_numel_list.append(tensor.numel())
         send_flatten_tensor.append(tensor.reshape(-1))
     for tensor in output_tensor_list:
-        recv_numel_list.append(tensor.size)
+        recv_numel_list.append(tensor.numel())
         recv_shape_list.append(tensor.shape)
 
     send_flatten_tensor = cat(send_flatten_tensor)
@@ -2183,10 +2584,7 @@ def all_to_all_single(output,
 
 def _check_tensor_list(tensor_list, tensor, group_size):
     """check all elements in tensor_list are type of Tensor or tuple or list"""
-    if not tensor_list or len(tensor_list) != group_size:
-        raise TypeError(
-            f"The argument list tensor len must be equal to group rank size, but got {len(tensor_list)}."
-        )
+    _check_group_tensor_list(tensor_list, group_size)
     if tensor.dtype != tensor_list[0].dtype:
         raise TypeError(
             f"The argument list tensor type must be equal to tensor type, but got {tensor_list[0].dtype}."
@@ -2197,13 +2595,17 @@ def _check_tensor_list(tensor_list, tensor, group_size):
         )
 
 
+def _check_group_tensor_list(tensor_list, group_size):
+    if not tensor_list or len(tensor_list) != group_size:
+        raise TypeError(
+            f"The argument list tensor len must be equal to group rank size, but got {len(tensor_list)}."
+        )
+
+
 def all_gather(tensor_list, tensor, group=None, async_op=False):
     """
     Gathers tensors from the specified communication group and returns the tensor list which is all gathered.
 
-    Note:
-        The tensors must have the same shape and format in all processes of the collection.
-
     Args:
         tensor_list (list[Tensor]): Output list.
         tensor (Tensor): The input tensor to be all gathered into tensor.
@@ -2258,7 +2660,7 @@ def all_gather(tensor_list, tensor, group=None, async_op=False):
 
     """
     _check_all_tensors(tensor_list)
-    _check_all_tensor_same_dtype_and_shape(tensor_list)
+    _check_all_tensor_same_dtype(tensor_list)
     if not isinstance(tensor, (Tensor, Tensor_)):
         raise TypeError("For all_gather_into_tensor, the input tensor must be tensor")
     if group is None:
@@ -2273,7 +2675,10 @@ def all_gather(tensor_list, tensor, group=None, async_op=False):
             f"The argument 'async_op' must be a bool, but got {type(async_op)}."
         )
     group_size = get_cache_group_size(group)
-    _check_tensor_list(tensor_list, tensor, group_size)
+    _check_group_tensor_list(tensor_list, group_size)
+    rank_id = get_group_rank_from_world_rank(get_rank(), group)
+    _check_output_shape(tensor, tensor_list[rank_id].shape, "all_gather")
+    _check_output_dtype(tensor, tensor_list[0].dtype, "all_gather")
     result = dist_comm_all_gather_op(tensor_list, tensor, group_size, group)
     _, handle = _deal_comm_outputs(result, async_op)
     return handle
@@ -2284,9 +2689,6 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=Fal
     Reduces and scatters tensors from the specified communication group and
     returns the tensor which is reduced and scattered.
 
-    Note:
-        The tensors must have the same shape and format in all processes of the collection.
-
     Args:
         output (Tensor): the output tensor.
         input_list (list[Tensor]): List of tensors to reduce and scatter.
@@ -2340,7 +2742,7 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=Fal
     """
 
     _check_all_tensors(input_list)
-    _check_all_tensor_same_dtype_and_shape(input_list)
+    _check_all_tensor_same_dtype(input_list)
     if not isinstance(output, (Tensor, Tensor_)):
         raise TypeError("For reduce_scatter, the output tensor must be tensor")
     if group is None:
@@ -2361,7 +2763,11 @@ def reduce_scatter(output, input_list, op=ReduceOp.SUM, group=None, async_op=Fal
             "For reduce_scatter, the input op value must be one of sum, prod, min, max"
         )
     rank_size = get_cache_group_size(group)
-    _check_tensor_list(input_list, output, rank_size)
+    _check_group_tensor_list(input_list, rank_size)
+
+    rank_id = get_group_rank_from_world_rank(get_rank(), group)
+    _check_output_shape(output, input_list[rank_id].shape, "reduce_scatter")
+    _check_output_dtype(output, input_list[0].dtype, "reduce_scatter")
     result = dist_comm_reduce_scatter_op(output, input_list, rank_size, op, group)
     _, handle = _deal_comm_outputs(result, async_op)
     return handle
@@ -2810,7 +3216,7 @@ def broadcast_object_list(object_list, src=0, group=None, device=None):
         offset = 0
         for i, item in enumerate(object_size):
             obj_size = item
-            obj_view = object_tensor[offset : offset + obj_size]
+            obj_view = object_tensor[offset: offset + obj_size]
             offset += obj_size
             object_list[i] = _tensor_to_object(obj_view, obj_size)
 

mindspore/mint/nn/__init__.py

@@ -1189,7 +1189,7 @@ class PixelShuffle(Cell):
         >>> pixel_shuffle = mint.nn.PixelShuffle(3)
         >>> input = mint.randn(1, 9, 4, 4)
         >>> output = pixel_shuffle(input)
-        >>> print(output.shape())
+        >>> print(output.shape)
             [1, 1, 12, 12]
     """