mindspore.communication.comm_func.all_to_all_with_output_shape
- mindspore.communication.comm_func.all_to_all_with_output_shape(output_shape_list, input_tensor_list, group=None)[source]
scatter and gather list of tensor to/from all rank according to input/output tensor list.
Note
tensor shape in output_shape_list and input_tensor_list should be match across ranks. Only support PyNative mode, Graph mode is not currently supported.
- Parameters
output_shape_list (Union[Tuple(Tensor), List(Tensor), Tuple(Tuple(int))]) – List of shape that indicate the gathered tensors shape from remote ranks.
input_tensor_list (Union[Tuple(Tensor), List(Tensor)]) – List of tensors to scatter to the remote rank.
group (str, optional) – The communication group to work on. Default: None, which means "hccl_world_group" on Ascend, "nccl_world_group" on GPU.
- Returns
Tuple(Tensor), the tensors gathered from remote ranks.
- Raises
- Supported Platforms:
Ascend
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 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 >>> >>> init() >>> this_rank = get_rank() >>> if this_rank == 0: >>> send_tensor_list = [Tensor(1.), Tensor([[2, 3], [4, 5.]])] >>> recv_tensor_list = [(), (2,)] >>> if this_rank == 1: >>> send_tensor_list = [Tensor([2, 2.]), Tensor([4, 5, 6, 7.])] >>> recv_tensor_list = [(2, 2), (4,)] >>> output = all_to_all_with_output_shape(recv_tensor_list, send_tensor_list) >>> print(output) rank 0: (Tensor(shape=[], dtype=Float32, value= 1), Tensor(shape=[2], dtype=Float32, value= [2.00000000e+00, 2.00000000e+00])) rank 1: (Tensor(shape=[2, 2], dtype=Float32, value= [[2.00000000e+00, 3.00000000e+00], [4.00000000e+00, 5.00000000e+00]]), Tensor(shape=[4], dtype=Float32, value=[4.00000000e+00, 5.00000000e+00, 6.00000000e+00, 7.00000000e+00]))