nvalchemi.distributed.DenseBatchPadder#
- class nvalchemi.distributed.DenseBatchPadder(*, nbmat_key='neighbor_matrix', initial_factor=1.15, grow_factor=1.30, stride=16)[source]#
Built-in
GraphPadderfor dense(N, K)neighbor-matrixBatches (AIMNet2).The batch-level counterpart of
DensePadder. The framework compiles the wholewrapper.forward, so the fixed-shape padding must land on the halo-paddedBatchbeforeadapt_inputruns — the same seamCOOPadderuses. This padder pads the atom-level storage group to a fixed atom capacity with inert dead atoms (zeros,Z=0, joined to the last graph) and repoints theneighbor_matrixsentinel so no real atom ever sees a dead atom as a neighbor.adapt_inputthen appends its own padding atom on top of this fixed-shape batch, so the compiled model input keeps a constant(n_cap + 1, …)shape across MD steps. The sentinel is repointed ton_cap— the index of that appended pad atom — so aimnet’scalc_masksmasks every dead / sentinel neighbor slot to zero.unpad()is a no-op: the owned-onlymol_sum(masked byn_owned) drops the dead rows from the energy, and the per-atom force output is sliced by the framework’s output consolidation. The framework restores the transient padded storage separately (the paddedBatchis reused in place across MD steps).- Parameters:
nbmat_key (str, default
"neighbor_matrix") – The dense neighbor-matrix node field whose padding sentinel (unused slots, set to the pre-pad node count) must be repointed to the appended pad-atom index.initial_factor (float) – Forwarded to
resolve_cap()for the"atoms"capacity.grow_factor (float) – Forwarded to
resolve_cap()for the"atoms"capacity.stride (int) – Forwarded to
resolve_cap()for the"atoms"capacity.