dgl.multiprocessing.pytorch.call_once_and_share

dgl.multiprocessing.pytorch.call_once_and_share(func, shape, dtype, rank=0)[source]

Invoke the function in a single process of the PyTorch distributed process group, and share the result with other processes.

Parameters:

func (callable) – Any callable that accepts no arguments and returns an arbitrary object.
shape (tuple[int]) – The shape of the shared tensor. Must match the output of func.
dtype (torch.dtype) – The data type of the shared tensor. Must match the output of func.
rank (int, optional) – The process ID to actually execute the function.