`nvalchemi.distributed`.OutputKind#

class nvalchemi.distributed.OutputKind(*values)[source]#

Per-output classification used by consolidation.

See module docstring for the design rationale.

Members#

PER_NODE: One row per atom. Halo storage: shape[0] == n_padded (owned + halo rows). Sharded storage: shape[0] == n_owned (owned only). Combine rule depends on MLIPSpec.owned_only_outputs / MLIPSpec.all_reduce_outputs membership and whether the key is in ModelConfig.autograd_outputs.
PER_GRAPH: One row per system. shape[0] == n_systems. Combine rule depends on autograd / all_reduce membership.
GLOBAL: Already globally-correct on every rank; passthrough. Rare — typically scalar metadata or replicated config tensors that come out of the wrapper unchanged.
UNKNOWN: Undeclared default. Consolidation falls back to the shape-based heuristic and logs a warning so the wrapper author knows to declare. Also accepted for non-tensor output values (which always pass through anyway).