collie.utils.broadcast_tensor

broadcast_tensor(tensor, dtype=None, src=0, shape=None, ndim=None, group=None)

从 src 广播 tensor。

该函数支持广播 tensor 的维度和类型。如果 ndim 和 shape 为 None 则会首先广播 tensor 的维度。

if rank == 0:
    logits = torch.ones(4,2)
else:
    logits = None

# 其它 rank 上并不知道 logits 的维度，因此 ndim 和 shape 均为 None
logits = broadcast_tensor(logits, src=0)

警告

请确保 ndim、shape 和 dtype 参数在所有 rank 上保持一致， 否则会导致分布式进程卡死。

参数:

tensor -- 要广播的张量。注意在 src rank 上必须为一个张量。