xnn.common.distributed.all_gather_cat#
- xnn.common.distributed.all_gather_cat(t, device)[source]#
Concatenate a flat tensor from every rank, in rank order.
Ranks may hold different lengths, including zero (a rank whose shard has no labelled structure for a target), so the lengths are exchanged first and each contribution is padded to the longest, gathered, and trimmed. The dtype travels with the lengths so an empty contribution can be typed like the others. Outside a distributed run the tensor is returned as is.
- Parameters:
t (torch.Tensor) – This rank’s one-dimensional contribution (any device).
device (torch.device) – This rank’s compute device, which selects where the collective runs.
- Returns:
The concatenation over ranks, on the CPU.
- Return type: