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, async_op=False)[源代码]

根据用户输入的张量列表，将对应的张量发送到远端设备，并从其他设备接收张量，返回一个接收的张量列表。

说明

各个设备之间发送和接收的张量形状需要互相匹配。 仅支持PyNative模式，目前不支持Graph模式。

参数：

  • output_shape_list (Union[Tuple(Tensor), List(Tensor), Tuple(Tuple(int))]) - 包含接收张量形状的列表。

  • input_tensor_list (Union[Tuple(Tensor), List(Tensor)]) - 包含发送到其他设备张量的列表。

  • group (str, 可选) - 通信所使用的通信组。默认值:None。为None时，在Ascend上将使用 hccl_world_group ，在GPU使用 nccl_world_group 。

  • async_op (bool, 可选) - 本算子是否是异步算子。默认值： False 。

返回：

Tuple(Tuple(Tensor), CommHandle), Tuple(Tensor)为从远端设备接收的张量列表。 若 async_op 是True，CommHandle是一个异步工作句柄。若 async_op 是False，CommHandle将返回None。

异常：

  • TypeError - input_tensor_list 中不全是张量类型。

  • TypeError - output_shape_list 中不全是张量或者元组类型。

  • TypeError - input_tensor_list 中张量的数据类型不全部一致。

支持平台：

Ascend

样例：

说明

运行以下样例之前，需要配置好通信环境变量。

针对Ascend/GPU/CPU设备，推荐使用msrun启动方式，无第三方以及配置文件依赖。详见 msrun启动 。

该样例需要在2卡环境下运行。

>>> import numpy as np
>>> import mindspore as ms
>>> import mindspore.communication as comm
>>>
>>> comm.init()
>>> this_rank = comm.get_rank()
>>> if this_rank == 0:
>>>     send_tensor_list = [ms.Tensor(1.), ms.Tensor([[2, 3], [4, 5.]])]
>>>     recv_tensor_list = [(), (2,)]
>>> if this_rank == 1:
>>>     send_tensor_list = [ms.Tensor([2, 2.]), ms.Tensor([4, 5, 6, 7.])]
>>>     recv_tensor_list = [(2, 2), (4,)]
>>> 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),
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]))