mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/ops/operations/comm_ops.py

@@ -52,12 +52,15 @@ class ReduceOp:
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -117,8 +120,9 @@ class AllReduce(Primitive):
 
     Args:
         op (str): Specifies an operation used for element-wise reductions, like sum, prod, max, and min.
-                  On the CPU, only 'sum' is supported. Default: ReduceOp.SUM.
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+                  On the CPU, only 'sum' is supported. Default: ``ReduceOp.SUM`` .
+        group (str): 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) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
@@ -139,14 +143,17 @@ class AllReduce(Primitive):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
 
-            This example should be run with multiple devices.
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
+
+            This example should be run with 2 devices.
 
         >>> import numpy as np
         >>> from mindspore.communication import init
@@ -170,6 +177,11 @@ class AllReduce(Primitive):
         >>> print(output)
         [[2. 2. 2. 2. 2. 2. 2. 2.]
          [2. 2. 2. 2. 2. 2. 2. 2.]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - AllReduce
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#allreduce>`_
+
     """
 
     @prim_attr_register
@@ -197,7 +209,8 @@ class AllGather(PrimitiveWithInfer):
         - Currently only supports GRAPH_MODE and it should be called in Cell.
 
     Args:
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+        group (str): 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) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
@@ -219,12 +232,15 @@ class AllGather(PrimitiveWithInfer):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with 2 devices.
 
@@ -253,6 +269,11 @@ class AllGather(PrimitiveWithInfer):
          [1. 1. 1. 1. 1. 1. 1. 1.]
          [1. 1. 1. 1. 1. 1. 1. 1.]
          [1. 1. 1. 1. 1. 1. 1. 1.]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - AllGather
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#allgather>`_
+
     """
 
     @prim_attr_register
@@ -268,9 +289,6 @@ class AllGather(PrimitiveWithInfer):
         self.add_prim_attr('mean_flag', False)
         self.add_prim_attr('no_eliminate', True)
 
-    def __call__(self, tensor):
-        raise NotImplementedError
-
     def infer_shape(self, x_shape):
         validator.check_positive_int(len(x_shape), "x shape", self.name)
         if x_shape[0] > 0:
@@ -288,8 +306,8 @@ class _MiniStepAllGather(PrimitiveWithInfer):
     internal use of parallel modules and cannot be called by users.
 
     Args:
-        group (str): The communication group to work on. Default: None.
-        grad_accumulation_step (int): The grad accumulation step. Default: None.
+        group (str): The communication group to work on. Default: ``None`` .
+        grad_accumulation_step (int): The grad accumulation step. Default: ``None`` .
     """
 
     @prim_attr_register
@@ -324,7 +342,7 @@ class _MicroStepAllGather(PrimitiveWithInfer):
     internal use of parallel modules and cannot be called by users.
 
     Args:
-        group (str): The communication group to work on. Default: None.
+        group (str): The communication group to work on. Default: ``None`` .
     """
 
     @prim_attr_register
@@ -364,7 +382,7 @@ class _HostAllGather(PrimitiveWithInfer):
         mpirun -output-filename log -merge-stderr-to-stdout -np 3 python test_host_all_gather.py
 
     Args:
-        group (Union[tuple[int],list[int]]): The rand_ids of communication group to work on. Default: None.
+        group (Union[tuple[int],list[int]]): The rand_ids of communication group to work on. Default: ``None`` .
 
     Raises:
         TypeError: If group is not a list nor tuple, or elements of group are not int.
@@ -410,16 +428,14 @@ class _HostAllGather(PrimitiveWithInfer):
 class ReduceScatter(Primitive):
     r"""
     Reduces and scatters tensors from the specified communication group.
-    For more details about it, please refer to `Distributed Set Communication Primitives - ReduceScatter \
-    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#reducescatter>`_ .
 
     Note:
         The tensors must have the same shape and format in all processes of the collection.
 
     Args:
-        op (str): Specifies an operation used for element-wise reductions,
-                  like SUM and MAX. Default: ReduceOp.SUM.
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+        op (str, optional): Specifies an operation used for element-wise reductions,
+                  like SUM and MAX. Default: ``ReduceOp.SUM`` .
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` .
 
     Inputs:
         - **input_x** (Tensor) - Input Tensor, suppose it has a shape :math:`(N, *)`, where `*`
@@ -441,12 +457,15 @@ class ReduceScatter(Primitive):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with 2 devices.
 
@@ -476,6 +495,11 @@ class ReduceScatter(Primitive):
          [2. 2. 2. 2. 2. 2. 2. 2.]
          [2. 2. 2. 2. 2. 2. 2. 2.]
          [2. 2. 2. 2. 2. 2. 2. 2.]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - ReduceScatter
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#reducescatter>`_
+
     """
 
     @prim_attr_register
@@ -490,9 +514,6 @@ class ReduceScatter(Primitive):
         self.add_prim_attr('fusion', 0)
         self.add_prim_attr('no_eliminate', True)
 
-    def __call__(self, tensor):
-        raise NotImplementedError
-
 
 class _HostReduceScatter(PrimitiveWithInfer):
     """
@@ -506,8 +527,8 @@ class _HostReduceScatter(PrimitiveWithInfer):
 
     Args:
         op (str): Specifies an operation used for element-wise reductions,
-                  like sum, max, avg. Default: ReduceOp.SUM.
-        group (Union[tuple[int],list[int]]): The rand_ids of communication group to work on. Default: None.
+                  like sum, max, avg. Default: ``ReduceOp.SUM`` .
+        group (Union[tuple[int],list[int]]): The rand_ids of communication group to work on. Default: ``None`` .
 
     Raises:
         TypeError: If op is not a string and group is not a list nor tuple,
@@ -558,7 +579,7 @@ class Broadcast(PrimitiveWithInfer):
     Args:
         root_rank (int): Source rank. Required in all processes except the one
                    that is sending the data.
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+        group (str, optional): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` .
 
     Inputs:
         - **input_x** (tuple[Tensor]) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
@@ -578,12 +599,15 @@ class Broadcast(PrimitiveWithInfer):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with multiple devices.
 
@@ -611,6 +635,11 @@ class Broadcast(PrimitiveWithInfer):
         (Tensor(shape[2,4], dtype=Int32, value=
         [[1, 1, 1, 1],
          [1, 1, 1, 1]]),)
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - Broadcast
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#broadcast>`_
+
     """
 
     @prim_attr_register
@@ -629,7 +658,7 @@ class Broadcast(PrimitiveWithInfer):
         if not isinstance(x_dtype, tuple):
             raise TypeError(f"For '{self.name}', the 'input_x' must be a tuple, but got {type(x_dtype).__name__}!")
         for _ele in x_dtype:
-            check_collective_target_dtype('input_x', _ele, self.name)
+            check_collective_target_dtype('tuple input_x', _ele, self.name)
         return x_dtype
 
 
@@ -671,7 +700,7 @@ class _AllSwap(PrimitiveWithCheck):
         self.add_prim_attr('order_enforce_skip', True)
 
     def __check__(self, tensor_in, send_size, recv_size):
-        validator.check_subclass("tensor_in", tensor_in['dtype'], mstype.tensor, self.name)
+        validator.check_subclass("tensor_in", tensor_in['dtype'], mstype.tensor_type, self.name)
         validator.check_tensor_dtype_valid("send_size", send_size['dtype'], [mstype.int64],
                                            self.name)
         validator.check_tensor_dtype_valid("recv_size", recv_size['dtype'], [mstype.int64],
@@ -699,11 +728,11 @@ class NeighborExchange(Primitive):
         The user needs to preset
         communication environment variables before running the following example, please check the details on the
         official website of `MindSpore \
-        <https://www.mindspore.cn/docs/en/r2.0/api_python/mindspore.ops.html#communication-operator>`_.
+        <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.ops.primitive.html#communication-operator>`_.
 
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
+        <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to.
@@ -711,7 +740,7 @@ class NeighborExchange(Primitive):
         recv_shapes (tuple(list(int))): Data shape which received from recv_rank_ids.
         send_shapes (tuple(list(int))): Data shape which send to the send_rank_ids.
         recv_type (type): Data type which received from recv_rank_ids
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+        group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` .
 
     Inputs:
         - **input_x** (tuple[Tensor]) - Shapes are same as args of send_shapes.
@@ -742,13 +771,18 @@ class NeighborExchange(Primitive):
         ...     def construct(self, x):
         ...         out = self.neighborexchange((x,))
         ...
-        >>> ms.set_context(mode=ms.GRAPH_MODE, device_target='Ascend')
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> init()
         >>> net = Net()
         >>> input_x = Tensor(np.ones([3, 3]), dtype = ms.float32)
         >>> output = net(input_x)
         >>> print(output)
         [[2. 2.], [2. 2.]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - NeighborExchange
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#neighborexchange>`_
+
     """
 
     @prim_attr_register
@@ -760,6 +794,7 @@ class NeighborExchange(Primitive):
         self.recv_shapes = recv_shapes
         self.send_shapes = send_shapes
         self.recv_type = recv_type
+        self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('no_eliminate', True)
 
     def __call__(self, tensor):
@@ -779,13 +814,13 @@ class AlltoAll(PrimitiveWithInfer):
     Note:
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
+        <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#notes>`_.
 
     Args:
         split_count (int): On each process, divide blocks into split_count number.
         split_dim (int): On each process, split blocks along the split_dim.
         concat_dim (int): On each process, gather the received blocks along the concat_dimension.
-        group (str): The communication group to work on. Default: "GlobalComm.WORLD_COMM_GROUP".
+        group (str): The communication group to work on. Default: ``GlobalComm.WORLD_COMM_GROUP`` .
 
     Inputs:
         - **input_x** (Tensor) - The shape of tensor is :math:`(x_1, x_2, ..., x_R)`.
@@ -809,12 +844,15 @@ class AlltoAll(PrimitiveWithInfer):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with 8 devices.
 
@@ -834,7 +872,7 @@ class AlltoAll(PrimitiveWithInfer):
         ...         out = self.alltoall(x)
         ...         return out
         ...
-        >>> ms.set_context(mode=ms.GRAPH_MODE, device_target='Ascend')
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> init()
         >>> net = Net()
         >>> rank_id = int(os.getenv("RANK_ID"))
@@ -842,6 +880,11 @@ class AlltoAll(PrimitiveWithInfer):
         >>> output = net(input_x)
         >>> print(output)
         [[[[0. 1. 2. 3. 4. 5. 6. 7.]]]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - AlltoAll
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#alltoall>`_
+
     """
 
     @prim_attr_register
@@ -882,15 +925,13 @@ class NeighborExchangeV2(Primitive):
     NeighborExchangeV2 is a collective communication operation.
 
     NeighborExchangeV2 sends data from the local rank to ranks in the `send_rank_ids`,
-    as while receive data from `recv_rank_ids`. Please refer to
-    `Distributed Set Communication Primitives - NeighborExchangeV2 \
-    <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#neighborexchangev2>`_
-    to learn about how the data is exchanged between neighborhood devices.
+    as while receive data from `recv_rank_ids`. Please refer to the tutorial examples
+    below to learn about how the data is exchanged between neighborhood devices.
 
     Note:
         This operator requires a full-mesh network topology, each device has the same vlan id, and the ip & mask are
         in the same subnet, please check the `details \
-        <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/communicate_ops.html#notes>`_.
+        <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#notes>`_.
 
     Args:
         send_rank_ids (list(int)): Ranks which the data is sent to. 8 rank_ids represents 8 directions, if one
@@ -902,8 +943,8 @@ class NeighborExchangeV2(Primitive):
         recv_lens (list(int)): Data lens which received from recv_rank_ids, 4 numbers represent the lens of
                                [recv_top, recv_bottom, recv_left, recv_right].
         data_format (str): Data format, only support NCHW now.
-        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.
+        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) - The Tensor before being exchanged. It has a shape of :math:`(N, C, H, W)`.
@@ -927,42 +968,68 @@ class NeighborExchangeV2(Primitive):
             Before running the following examples, you need to configure the communication environment variables.
 
             For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
-            Please see the `Ascend tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_ascend.html#preparations>`_
+            Please see the `rank table Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/rank_table.html>`_
             for more details.
 
-            For the GPU devices, users need to prepare the host file and mpi, please see the `GPU tutorial
-            <https://www.mindspore.cn/tutorials/experts/en/r2.0/parallel/train_gpu.html#preparation>`_ .
+            For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
+            <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/mpirun.html>`_ .
+
+            For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster
+            Startup <https://www.mindspore.cn/tutorials/experts/en/r2.2/parallel/dynamic_cluster.html>`_ .
 
             This example should be run with 2 devices.
 
         >>> import os
         >>> import mindspore as ms
-        >>> from mindspore import Tensor
         >>> from mindspore.communication import init
         >>> import mindspore.nn as nn
         >>> import mindspore.ops as ops
         >>> import numpy as np
-        >>> class Net(nn.Cell):
+        >>>
+        >>> class Net0(nn.Cell):
         ...     def __init__(self):
-        ...         super(Net, self).__init__()
-        ...         self.neighborexchangev2 = ops.NeighborExchangeV2(send_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1],
-        ...                                                          send_lens=[0, 1, 0, 0],
-        ...                                                          recv_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1],
-        ...                                                          recv_lens=[0, 1, 0, 0],
-        ...                                                          data_format="NCHW")
+        ...         super(Net0, self).__init__()
+        ...         self.neighbor_exchangev2 = ops.NeighborExchangeV2(send_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1],
+        ...                                                           send_lens=[0, 1, 0, 0],
+        ...                                                           recv_rank_ids=[-1, -1, -1, -1, 1, -1, -1, -1],
+        ...                                                           recv_lens=[0, 1, 0, 0], data_format="NCHW")
         ...
         ...     def construct(self, x):
-        ...         out = self.neighborexchangev2(x)
+        ...         out = self.neighbor_exchangev2(x)
         ...         return out
+        >>>
+        ... class Net1(nn.Cell):
+        ...     def __init__(self):
+        ...         super(Net1, self).__init__()
+        ...         self.neighbor_exchangev2 = ops.NeighborExchangeV2(send_rank_ids=[0, -1, -1, -1, -1, -1, -1, -1],
+        ...                                                           send_lens=[1, 0, 0, 0],
+        ...                                                           recv_rank_ids=[0, -1, -1, -1, -1, -1, -1, -1],
+        ...                                                           recv_lens=[1, 0, 0, 0], data_format="NCHW")
         ...
-        >>> ms.set_context(mode=ms.GRAPH_MODE, device_target='Ascend')
+        ...     def construct(self, x):
+        ...         out = self.neighbor_exchangev2(x)
+        ...         return out
+        >>>
+        >>> ms.set_context(mode=ms.GRAPH_MODE)
         >>> init()
-        >>> input_x = Tensor(np.ones([1, 1, 2, 2]), dtype = ms.float32)
-        >>> net = Net()
-        >>> output = net(input_x)
-        >>> print(output)
+        >>> rank_id = int(os.getenv("RANK_ID"))
+        >>> if (rank_id % 2 == 0):
+        >>>     input_x = ms.Tensor(np.ones([1, 1, 2, 2]), dtype = ms.float32)
+        >>>     net = Net0()
+        >>>     output = net(input_x)
+        >>>     print(output)
+        >>> else:
+        >>>     input_x = ms.Tensor(np.ones([1, 1, 2, 2]) * 2, dtype = ms.float32)
+        >>>     net = Net1()
+        >>>     output = net(input_x)
+        >>>     print(output)
         [[[[1. 1.], [1. 1.], [2. 2.]]]]
+
+    Tutorial Examples:
+        - `Distributed Set Communication Primitives - NeighborExchangeV2
+          <https://www.mindspore.cn/docs/en/r2.2/api_python/samples/ops/communicate_ops.html#neighborexchangev2>`_
+
     """
 
     @prim_attr_register
@@ -976,6 +1043,15 @@ class NeighborExchangeV2(Primitive):
         self.format = data_format
         self.add_prim_attr('group', _get_group(group))
         self.add_prim_attr('no_eliminate', True)
+        self.rank_size = get_group_size(_get_group(group))
+        for rank_id in send_rank_ids:
+            if rank_id != -1:
+                validator.check_number_range(rank_id, 0, self.rank_size, validator.INC_LEFT, int,
+                                             "rank_id in send_rank_ids")
+        for rank_id in recv_rank_ids:
+            if rank_id != -1:
+                validator.check_number_range(rank_id, 0, self.rank_size, validator.INC_LEFT, int,
+                                             "rank_id in recv_rank_ids")
 
     def __call__(self, tensor):
         raise NotImplementedError
@@ -987,9 +1063,9 @@ class _MirrorOperator(PrimitiveWithInfer):
     internal use of parallel modules and cannot be called by users.
 
     Args:
-        group (str): The communication group to work on. Default: None.
-        dev_num (int): The device number of the group. Default: None.
-        mean_flag (bool): Whether use mean in backward. Default: None.
+        group (str): The communication group to work on. Default: ``None`` .
+        dev_num (int): The device number of the group. Default: ``None`` .
+        mean_flag (bool): Whether use mean in backward. Default: ``None`` .
     """
 
     @prim_attr_register
@@ -1017,10 +1093,10 @@ class _MirrorMiniStepOperator(PrimitiveWithInfer):
     internal use of parallel modules and cannot be called by users.
 
     Args:
-        group (str): The communication group to work on. Default: None.
-        dev_num (int): The device number of the group. Default: None.
-        mean_flag (bool): Whether use mean in backward. Default: None.
-        grad_accumulation_step (int): The grad accumulation step. Default: None.
+        group (str): The communication group to work on. Default: ``None`` .
+        dev_num (int): The device number of the group. Default: ``None`` .
+        mean_flag (bool): Whether use mean in backward. Default: ``None`` .
+        grad_accumulation_step (int): The grad accumulation step. Default: ``None`` .
     """
 
     @prim_attr_register
@@ -1176,9 +1252,9 @@ class _MirrorMicroStepOperator(PrimitiveWithInfer):
     internal use of parallel modules and cannot be called by users.
 
     Args:
-        group (str): The communication group to work on. Default: None.
-        dev_num (int): The device number of the group. Default: None.
-        mean_flag (bool): Whether use mean in backward. Default: None.
+        group (str): The communication group to work on. Default: ``None`` .
+        dev_num (int): The device number of the group. Default: ``None`` .
+        mean_flag (bool): Whether use mean in backward. Default: ``None`` .
     """
 
     @prim_attr_register

mindspore/ops/operations/control_ops.py

@@ -25,8 +25,8 @@ class GeSwitch(PrimitiveWithInfer):
     """
     Adds control switch to data.
 
-    Switch data flows into false or true branch depending on the condition. If the condition is true,
-    the true branch will be activated, or vise verse.
+    Switch data flows into ``False`` or ``True`` branch depending on the condition. If the condition is ``True`` ,
+    the ``True`` branch will be activated, or vise verse.
 
     Inputs:
         - **data** (Union[Tensor, Number]) - The data to be used for switch control.
@@ -81,7 +81,7 @@ class GeSwitch(PrimitiveWithInfer):
 
     def infer_dtype(self, data_type, pred_type):
         validator.check_subclass(
-            "data", data_type, (mstype.tensor,) + mstype.number_type, self.name)
+            "data", data_type, (mstype.tensor_type,) + mstype.number_type, self.name)
         validator.check_tensor_dtype_valid("pred", pred_type, [mstype.bool_], self.name)
         return data_type, data_type