nvalchemi.distributed.scatter_to_owners#
- nvalchemi.distributed.scatter_to_owners(out)[source]#
Fold per-edge contributions written into ghost rows back to owners.
After a message-passing block scatters per-edge messages into nodes — leaving each rank’s partial sums in its ghost rows — this accumulates those partials into the owning ranks and re-broadcasts, so every rank’s owned and ghost rows hold the correct totals for the next block. Autograd-aware. Identity in single-process.
- Parameters:
out (torch.Tensor) –
(n_rows, *F)per-node tensor with this rank’s partial sums in the ghost rows.- Returns:
Same shape, with owners and ghosts carrying the cross-rank totals.
- Return type: