nvalchemi.distributed.refresh_neighbors

nvalchemi.distributed.refresh_neighbors(x)

Populate this rank’s neighbor (ghost) rows of a per-node tensor.

Call this at the start of a message-passing block that reads a node’s neighbors: it refreshes the ghost rows of x so each rank sees current neighbor features. Autograd-aware — gradients on the refreshed rows accumulate back to the owning ranks.

On the halo policy x is [owned | ghost (| dead padding)]; owned rows are exchanged into the ghost region and any trailing padding rows are preserved. In single-process this is the identity.

Parameters:

x (torch.Tensor) – (n_rows, *F) per-node features with this rank’s owned rows first.

Returns:

Same shape as x with the neighbor rows populated.

Return type:

torch.Tensor