nvalchemi.distributed.system_sum#

nvalchemi.distributed.system_sum(vals, idx, n, scope=Scope.OWNED)[source]#

Sum per-node values into per-system totals, without double-counting.

Each rank holds neighbor copies of atoms it does not own, so a plain scatter_add over all rows would over-count. This sums only this rank’s owned rows and (for Scope.OWNED) all-reduces across the mesh to the true global per-system total. In single-process it is a plain scatter_add over all rows.

Under compile it masks the ghost / dead rows by the routing’s n_owned tensor (a tensor mask, not a dynamic slice, so the partition can drift without forcing a recompile) and reduces over all rows. A wrapper calls it the same way in both modes.

Parameters:

vals (torch.Tensor) – (n_rows, *F) per-node values with owned rows first.
idx (torch.Tensor) – (n_rows,) integer system index for each row, in [0, n).
n (int) – Number of systems in the (global) batch.
scope (Scope, default Scope.OWNED) – OWNED → owned-only sum + cross-rank all-reduce (global total on every rank). LOCAL → this rank’s owned-only partial with no all-reduce (the framework’s output consolidation finishes it).

Returns:

(n, *F) per-system totals (replicated on every rank for OWNED; a per-rank partial for LOCAL).

Return type:

torch.Tensor