megengine.functional.distributed.remote_recv#

remote_recv(src_rank, device=None, inp=None)[source]#

Receive a tensor from another process.

Parameters:

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.

Return type:

Tensor

Returns:

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)