megengine.functional.distributed.remote_recv¶

remote_recv(src_rank, device=None, inp=None)[源代码]¶

Receive a tensor from another process.

参数

src_rank (int) – Rank of source process.
device (Optional[str]) – The specific device to execute this operator. None default device means the device of inp will be used. Specify “gpu0:1” to execute this operator on diffrent cuda stream, 1 is stream id, and default stream id is 0.
inp – Dummy input to determine received tensor type.

返回类型

Tensor

返回

Received tensor.

Examples:

if rank == 0:
    data = mge.tensor(1)
    # Tensor(1)
    F.distributed.remote_send(data, 1) # return None
else:
    data = F.distributed.remote_recv(0)
    # Tensor(1)

megengine.functional.distributed.remote_send

megengine.functional.distributed.all_gather