mindspore.ops.AlltoAllV

class mindspore.ops.AlltoAllV(group=None)

AllToAll which support uneven split.

Note

• Only support flatten tensor as input. input tensor should be flattened and concatenated before call this primitive.

Parameters

group (str) – The communication group to work on. Default: GlobalComm.WORLD_COMM_GROUP, which means "hccl_world_group" in Ascend.

Inputs:

• input_x (Tensor) - flatten tensor to scatter. The shape of tensor is \((x_1)\).

send_numel_list(Union[tuple[int], list[int], Tensor]): split numel to scatter to different remote rank. recv_numel_list(Union[tuple[int], list[int], Tensor]): split numel to gather from different remote rank.

Outputs:

Tensor. flattened and concatenated tensor gather from remote ranks. If gather result is empty, it will return a Tensor with value 0, which has no actual meaning.

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.

>>> from mindspore import ops
>>> import mindspore.nn as nn
>>> from mindspore.communication import init, get_rank
>>> from mindspore import Tensor
>>>
>>> init()
>>> rank = get_rank()
>>> class Net(nn.Cell):
...     def __init__(self):
...         super(Net, self).__init__()
...         self.all_to_all = ops.AlltoAllV()
...
...     def construct(self, x, send_numel_list, recv_numel_list):
...         return self.all_to_all(x, send_numel_list, recv_numel_list)
>>> send_numel_list = []
>>> recv_numel_list = []
>>> if rank == 0:
>>>    send_tensor = Tensor([0, 1, 2.])
>>>    send_numel_list = [1, 2]
>>>    recv_numel_list = [1, 2]
>>> elif rank == 1:
>>>    send_tensor = Tensor([3, 4, 5.])
>>>    send_numel_list = [2, 1]
>>>    recv_numel_list = [2, 1]
>>> net = Net()
>>> output = net(send_tensor, send_numel_list, recv_numel_list)
>>> print(output)
rank 0:
[0. 3. 4]
rank 1:
[1. 2. 5]