mindspore 2.4.1__cp310-cp310-manylinux1_x86_64.whl → 2.5.0__cp310-cp310-manylinux1_x86_64.whl

mindspore/communication/comm_func.py

@@ -28,7 +28,9 @@ from mindspore.ops.auto_generate.gen_ops_prim import (inner_comm_all_reduce_op,
                                                       inner_comm_all_to_all_v_op, inner_comm_irecv_op,
                                                       inner_comm_isend_op, inner_comm_reduce_scatter_op)
 from mindspore._c_expression import CommHandle as CommHandle_
+from mindspore._c_expression.typing import Type
 from mindspore import jit_class
+from mindspore.common.api import _pynative_executor
 
 __all__ = [
     'all_reduce',
@@ -218,19 +220,18 @@ def all_reduce(tensor, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP, async
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
         >>> import numpy as np
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import all_reduce
-        >>> from mindspore import Tensor
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
         >>>
-        >>> init()
-        >>> input_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
-        >>> output = all_reduce(input_tensor)
+        >>> comm.init()
+        >>> input_tensor = ms.Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> output, _ = comm.comm_func.all_reduce(input_tensor)
         >>> print(output)
         [[2. 2. 2. 2. 2. 2. 2. 2.]
          [2. 2. 2. 2. 2. 2. 2. 2.]]
@@ -284,22 +285,18 @@ def all_gather_into_tensor(tensor, group=GlobalComm.WORLD_COMM_GROUP, async_op=F
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/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.communication import init
-        >>> from mindspore.communication.comm_func import all_gather_into_tensor
-        >>> from mindspore import Tensor
+        >>> import mindspore.communication as comm
         >>>
-        >>> ms.set_context(mode=ms.GRAPH_MODE)
-        >>> init()
-        >>> input_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
-        >>> output = all_gather_into_tensor(input_tensor)
+        >>> comm.init()
+        >>> input_tensor = ms.Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> output, _ = comm.comm_func.all_gather_into_tensor(input_tensor)
         >>> print(output)
         [[1. 1. 1. 1. 1. 1. 1. 1.]
          [1. 1. 1. 1. 1. 1. 1. 1.]
@@ -358,21 +355,18 @@ def reduce_scatter_tensor(tensor, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_G
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/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
-        >>> from mindspore.communication.comm_func import reduce_scatter_tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
         >>>
-        >>> ms.set_context(mode=ms.GRAPH_MODE)
-        >>> init()
-        >>> input_tensor = Tensor(np.ones([8, 8]).astype(np.float32))
-        >>> output = reduce_scatter_tensor(input_tensor)
+        >>> comm.init()
+        >>> input_tensor = ms.Tensor(np.ones([8, 8]).astype(np.float32))
+        >>> output, _ = comm.comm_func.reduce_scatter_tensor(input_tensor)
         >>> print(output)
         [[2. 2. 2. 2. 2. 2. 2. 2.]
          [2. 2. 2. 2. 2. 2. 2. 2.]
@@ -430,22 +424,20 @@ def reduce(tensor, dst, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP):
             without any third-party or configuration file dependencies.
 
             Please see the `msrun start up
-            <https://www.mindspore.cn/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 4 devices.
 
-        >>> from mindspore import ops
-        >>> import mindspore.nn as nn
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import reduce
-        >>> from mindspore import Tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
+        >>>
         >>> # Launch 4 processes.
-        >>> init()
+        >>> comm.init()
         >>> dest_rank=1
-        >>> input_tensor = Tensor(np.ones([2, 8]).astype(np.float32))
-        >>> output = reduce(input_tensor)
+        >>> input_tensor = ms.Tensor(np.ones([2, 8]).astype(np.float32))
+        >>> output = comm.comm_func.reduce(input_tensor, dst=dest_rank)
         >>> print(output)
         Process with rank 1: [[4. 4. 4. 4. 4. 4. 4. 4.]
                              [4. 4. 4. 4. 4. 4. 4. 4.]],
@@ -494,27 +486,36 @@ class P2POp:
 
     Examples:
         >>> import numpy as np
-        >>> import mindspore
-        >>> from mindspore.communication.comm_func import P2POp, isend, irecv
-        >>> from mindspore import Tensor
-        >>> send_tensor = Tensor(1.)
-        >>> send_op = P2POp('isend', send_tensor, 1)
-        >>> send_op = P2POp(isend, send_tensor, 1)
-        >>> recv_tensor = Tensor(0.)
-        >>> recv_op = P2POp('irecv', recv_tensor, 0)
-        >>> recv_op = P2POp(irecv, recv_tensor, 0)
-        >>> recv_op = P2POp('irecv', (), 0, recv_dtype=mindspore.float32)
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
+        >>>
+        >>> send_tensor = ms.Tensor(1.)
+        >>> send_op = comm.comm_func.P2POp('isend', send_tensor, 1)
+        >>> send_op = comm.comm_func.P2POp(comm.comm_func.isend, send_tensor, 1)
+        >>> recv_tensor = ms.Tensor(0.)
+        >>> recv_op = comm.comm_func.P2POp('irecv', recv_tensor, 0)
+        >>> recv_op = comm.comm_func.P2POp(comm.comm_func.irecv, recv_tensor, 0)
+        >>> recv_op = comm.comm_func.P2POp('irecv', (), 0, recv_dtype=ms.float32)
     """
 
     def __init__(self, op, tensor, peer, group=None, tag=0, *, recv_dtype=None):
         self.op = op
         self.tensor = tensor
+        if not isinstance(peer, int):
+            raise TypeError(f"peer must be type of int, but got type of {type(peer)}")
+
+        if recv_dtype and not isinstance(recv_dtype, Type):
+            raise TypeError(f"recv_dtype must be type of mindspore dtype, but got type of {type(recv_dtype)}")
+
         self.peer = peer
         self.group = group
         self.tag = tag
         self.recv_dtype = recv_dtype
 
     def __new__(cls, op, tensor, peer, group=None, tag=0, recv_dtype=None):
+        if not (isinstance(op, str) or callable(op)):
+            raise TypeError(f"op must be type of string or function, but got type of {type(op)}")
+
         if isinstance(op, str):
             op_name = op
         else:
@@ -560,31 +561,29 @@ def batch_isend_irecv(p2p_op_list):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
         >>> import numpy as np
-        >>> import mindspore
-        >>> from mindspore.communication import init, get_rank, get_group_size
-        >>> from mindspore.communication.comm_func import batch_isend_irecv, P2POp
-        >>> from mindspore import Tensor
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
         >>>
-        >>> init()
-        >>> this_rank = get_rank()
-        >>> world_size = get_group_size()
+        >>> comm.init()
+        >>> this_rank = comm.get_rank()
+        >>> world_size = comm.get_group_size()
         >>> next_rank = (this_rank + 1) % world_size
         >>> prev_rank = (this_rank + world_size - 1) % world_size
         >>>
-        >>> send_tensor = Tensor(this_rank + 1, dtype=mindspore.float32)
-        >>> recv_tensor = Tensor(0., dtype=mindspore.float32)
+        >>> send_tensor = ms.Tensor(this_rank + 1, dtype=ms.float32)
+        >>> recv_tensor = ms.Tensor(0., dtype=ms.float32)
         >>>
-        >>> send_op = P2POp('isend', send_tensor, next_rank)
-        >>> recv_op = P2POp('irecv', recv_tensor, prev_rank)
+        >>> send_op = comm.comm_func.P2POp('isend', send_tensor, next_rank)
+        >>> recv_op = comm.comm_func.P2POp('irecv', recv_tensor, prev_rank)
         >>>
         >>> p2p_op_list = [send_op, recv_op]
-        >>> output = batch_isend_irecv(p2p_op_list)
+        >>> output = comm.comm_func.batch_isend_irecv(p2p_op_list)
         >>> print(output)
         rank 0:
         (Tensor(shape=[], dtype=Float32, value= 0), Tensor(shape=[], dtype=Float32, value= 2))
@@ -597,6 +596,10 @@ def batch_isend_irecv(p2p_op_list):
     receive_shapes = []
     receive_dtypes = []
     tags = []
+
+    if not isinstance(p2p_op_list, list):
+        raise TypeError(f"p2p_op_list must be type of list, but got type of {p2p_op_list}.")
+
     if not p2p_op_list:
         raise TypeError(f"p2p_op_list can not be empty list.")
     group = p2p_op_list[0].group
@@ -676,20 +679,20 @@ def scatter_tensor(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
-        >>> import mindspore as ms
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import scatter_tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
+        >>>
         >>> # Launch 2 processes.
         >>>
-        >>> init()
+        >>> comm.init()
         >>> input = ms.Tensor(np.arange(8).reshape([4, 2]).astype(np.float32))
-        >>> out = scatter_tensor(tensor=data, src=0)
+        >>> out = comm.comm_func.scatter_tensor(tensor=input, src=0)
         >>> print(out)
         # rank_0
         [[0. 1.]
@@ -741,22 +744,20 @@ def gather_into_tensor(tensor, dst=0, group=GlobalComm.WORLD_COMM_GROUP):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
         >>> import numpy as np
         >>> import mindspore as ms
-        >>> import mindspore.nn as nn
-        >>> from mindspore.communication import init
-        >>> from mindspore import Tensor
-        >>> from mindspore.communication.comm_func import gather_into_tensor
+        >>> import mindspore.communication as comm
+        >>>
         >>> # Launch 2 processes.
         >>>
-        >>> init()
-        >>> input = Tensor(np.arange(4).reshape([2, 2]).astype(np.float32))
-        >>> output = gather_into_tensor(tensor=data, dst=0)
+        >>> comm.init()
+        >>> input = ms.Tensor(np.arange(4).reshape([2, 2]).astype(np.float32))
+        >>> output = comm.comm_func.gather_into_tensor(tensor=input, dst=0)
         >>> print(output)
         Process with rank 0: [[0. 1.],
                               [2. 3.],
@@ -804,21 +805,20 @@ def broadcast(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/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
-        >>> from mindspore.communication.comm_func import broadcast
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
+        >>>
         >>> # Launch 2 processes.
         >>>
-        >>> init()
+        >>> comm.init()
         >>> data = ms.Tensor(np.arange(8).reshape([2, 4]).astype(np.float32))
-        >>> out = broadcast(tensor=data, src=0)
+        >>> out = comm.comm_func.broadcast(tensor=data, src=0)
         [[0. 1. 2. 3.]
          [4. 5. 6. 7.]]
 
@@ -858,29 +858,37 @@ def barrier(group=GlobalComm.WORLD_COMM_GROUP):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import barrier
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
+        >>>
         >>> # Launch 2 processes.
-        >>> init()
-        >>> barrier()
+        >>> comm.init()
+        >>> comm.comm_func.barrier()
 
     Tutorial Examples:
         - `Distributed Set Communication Primitives - Barrier
           <https://www.mindspore.cn/docs/en/master/api_python/samples/ops/communicate_ops.html#barrier>`_
     """
+    if not isinstance(group, str):
+        raise TypeError(f"group must be type of string, but got {type(group)}")
     _op = _get_cache_prim(P.Barrier)(group)
     return _op()
 
 
-def _deal_comm_outputs(output, async_op):
+def _deal_comm_outputs(output, async_op, exec_sync=False):
+    """
+    deal with comm ops outputs.
+    """
     if isinstance(output, tuple):
         if not async_op:
             output[1].wait()
+            if exec_sync:
+                _pynative_executor.sync()
             return (output[0], None)
         return output
 
@@ -918,21 +926,34 @@ def send(tensor, dst=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
-        >>> from mindspore import ops
-        >>> import mindspore.nn as nn
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import send
-        >>> from mindspore import Tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore.communication import init
+        >>> from mindspore.communication.comm_func import send, recv
+        >>> from mindspore.communication import get_rank, get_group_size
         >>>
+        >>> np.random.seed(1)
         >>> init()
-        >>> input_ = Tensor(np.ones([2, 8]).astype(np.float32))
-        >>> send(input_, 0)
+        >>> rank = get_rank()
+        >>> size = get_group_size()
+        >>> x = np.ones([2, 2]).astype(np.float32) * 0.01 * (rank + 1)
+        >>> x2 = np.ones([2, 2]).astype(np.float32)
+        >>>
+        >>>
+        >>> if rank < size / 2:
+        >>>     _x = ms.Tensor(x)
+        >>>     send(_x, rank + size // 2)
+        >>> else:
+        >>>     _x2 = ms.Tensor(x2)
+        >>>     output = recv(_x2, rank - size // 2)
+        >>>     print(output)
+        [[0.01  0.01]
+         [0.01  0.01]]
     """
     if not isinstance(tensor, (Tensor, Tensor_)):
         raise TypeError("For send, the input tensor must be tensor")
@@ -979,29 +1000,34 @@ def recv(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
-        >>> from mindspore import ops
-        >>> import mindspore.nn as nn
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import recv
-        >>> from mindspore import Tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore.communication import init
+        >>> from mindspore.communication.comm_func import send, recv
+        >>> from mindspore.communication import get_rank, get_group_size
         >>>
-        # Launch 2 processes.
-        Process 0 send the following array to Process 1
-        [[ 0.  1.]
-         [ 2.  3.]]
+        >>> np.random.seed(1)
         >>> init()
-        >>> x = ms.Tensor(np.zeros([2, 2]))
-        # Process 1 receive tensor from Process 0.
-        >>> out = recv(x, src=0)
-        >>> print(out)
-        [[ 0.  1.]
-         [ 2.  3.]]
+        >>> rank = get_rank()
+        >>> size = get_group_size()
+        >>> x = np.ones([2, 2]).astype(np.float32) * 0.01 * (rank + 1)
+        >>> x2 = np.ones([2, 2]).astype(np.float32)
+        >>>
+        >>>
+        >>> if rank < size / 2:
+        >>>     _x = ms.Tensor(x)
+        >>>     send(_x, rank + size // 2)
+        >>> else:
+        >>>     _x2 = ms.Tensor(x2)
+        >>>     output = recv(_x2, rank - size // 2)
+        >>>     print(output)
+        [[0.01  0.01]
+         [0.01  0.01]]
     """
     if not isinstance(tensor, (Tensor, Tensor_)):
         raise TypeError("For recv, the input tensor must be tensor")
@@ -1049,22 +1075,35 @@ def isend(tensor, dst=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
-        >>> from mindspore import ops
-        >>> import mindspore.nn as nn
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import isend
-        >>> from mindspore import Tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore.communication import init
+        >>> from mindspore.communication.comm_func import isend, irecv
+        >>> from mindspore.communication import get_rank, get_group_size
         >>>
+        >>> np.random.seed(1)
         >>> init()
-        >>> input_ = Tensor(np.ones([2, 8]).astype(np.float32))
-        >>> handle = isend(input_, 0)
-        >>> handle.wait()
+        >>> rank = get_rank()
+        >>> size = get_group_size()
+        >>> x = np.ones([2, 2]).astype(np.float32) * 0.01 * (rank + 1)
+        >>> x2 = np.ones([2, 2]).astype(np.float32)
+        >>>
+        >>>
+        >>> if rank < size / 2:
+        >>>     _x = ms.Tensor(x)
+        >>>     isend(_x, rank + size // 2)
+        >>> else:
+        >>>     _x2 = ms.Tensor(x2)
+        >>>     output, handle = irecv(_x2, rank - size // 2)
+        >>>     handle.wait()
+        >>>     print(output)
+        [[0.01  0.01]
+         [0.01  0.01]]
     """
     if not isinstance(tensor, (Tensor, Tensor_)):
         raise TypeError("For isend, the input tensor must be tensor")
@@ -1114,30 +1153,35 @@ def irecv(tensor, src=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
-        >>> from mindspore import ops
-        >>> import mindspore.nn as nn
-        >>> from mindspore.communication import init
-        >>> from mindspore.communication.comm_func import irecv
-        >>> from mindspore import Tensor
         >>> import numpy as np
+        >>> import mindspore as ms
+        >>> from mindspore.communication import init
+        >>> from mindspore.communication.comm_func import isend, irecv
+        >>> from mindspore.communication import get_rank, get_group_size
         >>>
-        # Launch 2 processes.
-        Process 0 send the following array to Process 1
-        [[ 0.  1.]
-         [ 2.  3.]]
+        >>> np.random.seed(1)
         >>> init()
-        >>> x = ms.Tensor(np.zeros([2, 2]))
-        # Process 1 receive tensor from Process 0.
-        >>> out, handle = irecv(x, src=0)
-        >>> handle.wait()
-        >>> print(out)
-        [[ 0.  1.]
-         [ 2.  3.]]
+        >>> rank = get_rank()
+        >>> size = get_group_size()
+        >>> x = np.ones([2, 2]).astype(np.float32) * 0.01 * (rank + 1)
+        >>> x2 = np.ones([2, 2]).astype(np.float32)
+        >>>
+        >>>
+        >>> if rank < size / 2:
+        >>>     _x = ms.Tensor(x)
+        >>>     isend(_x, rank + size // 2)
+        >>> else:
+        >>>     _x2 = ms.Tensor(x2)
+        >>>     output, handle = irecv(_x2, rank - size // 2)
+        >>>     handle.wait()
+        >>>     print(output)
+        [[0.01  0.01]
+         [0.01  0.01]]
     """
     group = _get_group(group)
     _src = _get_group_rank_from_world_rank_from_cache_helper(src, group)
@@ -1185,27 +1229,24 @@ def all_to_all_with_output_shape(output_shape_list, input_tensor_list, group=Non
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
         >>> import numpy as np
-        >>> import mindspore
-        >>> from mindspore.communication import init, get_rank, get_group_size
-        >>> from mindspore.communication.comm_func import all_to_all_with_output_shape
-        >>> from mindspore import Tensor
-        >>> from mindspore.ops import zeros
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
         >>>
-        >>> init()
-        >>> this_rank = get_rank()
+        >>> comm.init()
+        >>> this_rank = comm.get_rank()
         >>> if this_rank == 0:
-        >>>     send_tensor_list = [Tensor(1.), Tensor([[2, 3], [4, 5.]])]
+        >>>     send_tensor_list = [ms.Tensor(1.), ms.Tensor([[2, 3], [4, 5.]])]
         >>>     recv_tensor_list = [(), (2,)]
         >>> if this_rank == 1:
-        >>>     send_tensor_list = [Tensor([2, 2.]), Tensor([4, 5, 6, 7.])]
+        >>>     send_tensor_list = [ms.Tensor([2, 2.]), ms.Tensor([4, 5, 6, 7.])]
         >>>     recv_tensor_list = [(2, 2), (4,)]
-        >>> output = all_to_all_with_output_shape(recv_tensor_list, send_tensor_list)
+        >>> output, _ = comm.comm_func.all_to_all_with_output_shape(recv_tensor_list, send_tensor_list)
         >>> print(output)
         rank 0:
         (Tensor(shape=[], dtype=Float32, value= 1),
@@ -1256,17 +1297,17 @@ def all_to_all_with_output_shape(output_shape_list, input_tensor_list, group=Non
     return (tuple(result), handle)
 
 
-def _get_all_to_all_single_numel_list(tensor, output_shape, output_split_sizes, input_split_sizes, group):
+def _get_all_to_all_single_numel_list(tensor_shape, output_shape, output_split_sizes, input_split_sizes, group):
     """get numel list for all_to_all_single."""
     global _GROPU_SIZE_CACHE
     if _is_split_sizes_empty(input_split_sizes):
         if group not in _GROPU_SIZE_CACHE:
             _GROPU_SIZE_CACHE[group] = get_group_size(group)
         _world_size = _GROPU_SIZE_CACHE[group]
-        if tensor.shape[0] % _world_size != 0:
+        if tensor_shape[0] % _world_size != 0:
             raise ValueError("input shape at dim 0 must be divided by world_size, "
-                             f"but got {tensor.shape[0]} and {_world_size}.")
-        _split_size = tensor.shape[0] // _world_size
+                             f"but got {tensor_shape[0]} and {_world_size}.")
+        _split_size = tensor_shape[0] // _world_size
         input_split_sizes = (_split_size,) * _world_size
     if _is_split_sizes_empty(output_split_sizes):
         if group not in _GROPU_SIZE_CACHE:
@@ -1283,7 +1324,7 @@ def _get_all_to_all_single_numel_list(tensor, output_shape, output_split_sizes,
         _split_size = shape_dim_0 // _world_size
         output_split_sizes = (_split_size,) * _world_size
 
-    send_size_without_first_dim = _get_size(tensor.shape[1:])
+    send_size_without_first_dim = _get_size(tensor_shape[1:])
     send_numel_list = [size * send_size_without_first_dim for size in input_split_sizes]
 
     recv_size_without_first_dim = None
@@ -1298,6 +1339,9 @@ def _get_all_to_all_single_numel_list(tensor, output_shape, output_split_sizes,
     return send_numel_list, recv_numel_list, recv_shape_without_first_dim
 
 
+_ALL_TO_ALL_CACHE = {}
+
+
 def all_to_all_single_with_output_shape(output_shape, tensor, output_split_sizes=None,
                                         input_split_sizes=None, group=None, async_op=False):
     """
@@ -1339,36 +1383,25 @@ def all_to_all_single_with_output_shape(output_shape, tensor, output_split_sizes
             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/docs/zh-CN/master/model_train/parallel/msrun_launcher.html>`_
+            <https://www.mindspore.cn/docs/en/master/model_train/parallel/msrun_launcher.html>`_
             for more details.
 
             This example should be run with 2 devices.
 
         >>> import numpy as np
-        >>> import mindspore
-        >>> from mindspore.communication import init, get_rank, get_group_size
-        >>> from mindspore.communication.comm_func import all_to_all_single_with_output_shape
-        >>> from mindspore import Tensor
-        >>> from mindspore.ops import zeros
+        >>> import mindspore as ms
+        >>> import mindspore.communication as comm
         >>>
-        >>> init()
-        >>> this_rank = get_rank()
-        >>> if this_rank == 0:
-        >>>     output_shape = (3, 3)
-        >>>     tensor = Tensor([[0, 1, 2.], [3, 4, 5], [6, 7, 8]])
-        >>>     result = all_to_all_single_with_output_shape(output_shape, tensor, [2, 1], [2, 1])
-        >>> if this_rank == 1:
-        >>>     output_shape = (2, 3)
-        >>>     tensor = Tensor([[9, 10., 11], [12, 13, 14]])
-        >>>     result = all_to_all_single_with_output_shape(output_shape, tensor)
+        >>> comm.init()
+        >>> rank = comm.get_rank()
+        >>> input = ms.Tensor([0, 1]) + rank * 2
+        >>> output_shape = (2,)
+        >>> result, _ = comm.comm_func.all_to_all_single_with_output_shape(output_shape, input)
         >>> print(result)
         rank 0:
-        [[ 0.  1.  2.]
-         [ 3.  4.  5.]
-         [ 9. 10. 11.]]
+        [ 0.  2.]
         rank 1:
-        [[ 6.  7.  8.]
-         [12. 13. 14.]]
+        [ 1.  3.]
 
     """
 
@@ -1378,8 +1411,17 @@ def all_to_all_single_with_output_shape(output_shape, tensor, output_split_sizes
         group = GlobalComm.WORLD_COMM_GROUP
 
     split_sizes_empty = _is_split_sizes_empty(output_split_sizes) and _is_split_sizes_empty(input_split_sizes)
-    send_numel_list, recv_numel_list, recv_shape_without_first_dim = \
-        _get_all_to_all_single_numel_list(tensor, output_shape, output_split_sizes, input_split_sizes, group)
+    if isinstance(output_split_sizes, list):
+        output_split_sizes = tuple(output_split_sizes)
+    if isinstance(input_split_sizes, list):
+        input_split_sizes = tuple(input_split_sizes)
+    global _ALL_TO_ALL_CACHE
+    tensor_shape = output_shape
+    cache_key = (tensor_shape, output_shape, output_split_sizes, input_split_sizes, group)
+    if cache_key not in _ALL_TO_ALL_CACHE:
+        _ALL_TO_ALL_CACHE[cache_key] = _get_all_to_all_single_numel_list(*cache_key)
+    send_numel_list, recv_numel_list, recv_shape_without_first_dim = _ALL_TO_ALL_CACHE[cache_key]
+
     tensor = _contiguous(tensor)
     _input = tensor.reshape(-1)
     group = GlobalComm.WORLD_COMM_GROUP if group is None else _get_group(group)