nvalchemi.distributed.DensePadder

class nvalchemi.distributed.DensePadder(*, count_key, nbmat_key, row_pads, atom_output_keys=(), n_systems_key=None, cap_key='atoms', initial_factor=1.15, grow_factor=1.15, stride=16)

Built-in GraphPadder for dense (N, K) neighbor-matrix graphs.

The dense counterpart of COOPadder: the graph is an (N, K) neighbor matrix that rides the atom dimension (no separate edge dim). Pads the per-atom row fields to a fixed atom capacity and repoints the neighbor matrix’s padding sentinel to an isolated dead atom (the last row); unpad slices the dead rows off the per-atom outputs.

Parametrized by the model’s field names: count_key (field whose row count is the atom count), nbmat_key (the neighbor matrix), row_pads (per-atom field -> pad fill value; pass LAST_SYSTEM to pad a per-system label with the last system’s index), and atom_output_keys (per-atom outputs that get dead rows stripped in unpad()).

Layout assumption: the input’s last pre-pad row is the model’s own padding/sentinel atom, so the real atom count is n_rows - 1 and neighbor entries >= n_rows - 1 are the sentinel — both get repointed to the dead row.

Parameters:

• count_key (str)

• nbmat_key (str)

• row_pads (dict[str, Any])

• atom_output_keys (tuple[str, ...])

• n_systems_key (str | None)

• cap_key (str)

• initial_factor (float)

• grow_factor (float)

• stride (int)

pad(data, cap_state, cap_atoms=True)

Resolve the atom cap from cap_state and pad the dense fields.

data is the model’s plain-tensor input dict; returns a shallow copy with the row fields + neighbor matrix padded to the atom cap.

Parameters:

• data (dict[str, Any])

• cap_state (dict[str, int])

• cap_atoms (bool)

Return type:

dict[str, Any]

unpad(output, n_real=None)

Slice the dead-atom rows off the per-atom outputs.

Strips to n_real when given (eager / sharded paths, where no pad() ran), else to the count the matching pad() stashed.

Parameters:

• output (dict[str, Any])

• n_real (int | None)

Return type:

dict[str, Any]