nvalchemi.distributed.GraphPadder

class nvalchemi.distributed.GraphPadder(*args, **kwargs)

How a model’s graph representation is padded to a fixed capacity.

Declared on CompilePolicy.graph_padder; the default is COOPadder. pad() runs before the compiled forward and unpad() on the raw output, so the wrapper stays distribution-agnostic and the compiled graph shape stays constant.

The padder owns capacity resolution: it is handed a mutable cap_state dict (persistent across MD steps) and sizes its own caps with resolve_cap(). This matters because when a capacity becomes knowable is representation-specific — a COO graph knows its atom/edge counts up front, but a model that rebuilds its graph inside pad() only learns its edge count partway through.

Padding must be inert: dead atoms carry no contribution (e.g. Z=0, parked beyond any cutoff) and dead edges have a zero envelope (e.g. a self-loop longer than the cutoff, or a non-degenerate image so spherical harmonics don’t NaN). The owned-only output consolidation drops the dead rows regardless, but they must not perturb the real atoms’ values.

pad(data, cap_state, cap_atoms=True)

Return data padded to fixed per-rank capacities.

Resolve the capacities with resolve_cap() against cap_state (grow-only, persistent across steps), then pad data and return it.

cap_atoms (strategy-supplied) selects whether the atom dim is capped too or only edges: halo caps both (owned+ghost fluctuate); a graph-parallel node partition caps edges only (fixed atom set). Padders whose layout only ever caps atoms may ignore it.

Parameters:

• data (Any)

• cap_state (dict[str, int])

• cap_atoms (bool)

Return type:

Any

unpad(output, n_real=None)

Drop the dead-atom / dead-edge rows from a raw model output.

n_real overrides the real row count to strip to; pass it on paths where no pad() ran (e.g. eager / sharded), else leave it None to use the count the matching pad() recorded.

Parameters:

• output (Any)

• n_real (int | None)

Return type:

Any