nvalchemi.distributed.OutputSpec

class nvalchemi.distributed.OutputSpec(kind=OutputKind.UNKNOWN, reduce=Reduce.NONE)

How one named model output is shaped and combined under DD.

The single per-output declaration that collapses the three parallel sets output_kinds / all_reduce_outputs / owned_only_outputs into one place:

outputs={"stress": OutputSpec(kind=OutputKind.PER_GRAPH,
                              reduce=Reduce.ALL_REDUCE)}

MLIPSpec accepts outputs={name: OutputSpec} and lowers it onto those legacy fields, so consolidation and serialization are unchanged.

Parameters:

• kind (OutputKind)

• reduce (Reduce)