nvalchemi.distributed.MLIPSpec

class nvalchemi.distributed.MLIPSpec(distribution, owned_only_outputs=<factory>, all_reduce_outputs=<factory>, output_kinds=<factory>, system_reductions=True, node_energy_key=None, node_virial_key=None, gp_replicate_geometry=False, outputs=None, compile=None)

What an MLIP needs from the distributed framework.

Wraps a DistributionSpec and adds output-classification sets keyed by output name ("forces", "stress", etc.). The framework reads it via BaseModelMixin.distribution_spec.

The recommended construction declares each output once via outputs=:

MLIPSpec(
    distribution=DistributionSpec(
        policy=HaloStoragePolicy(),
        custom_ops=(...),
    ),
    outputs={
        "energy": OutputSpec(OutputKind.PER_GRAPH),
        "forces": OutputSpec(OutputKind.PER_NODE),
        "stress": OutputSpec(OutputKind.PER_GRAPH, Reduce.ALL_REDUCE),
    },
)

outputs is lowered in __post_init__ onto the canonical fields (output_kinds / owned_only_outputs / all_reduce_outputs), which are what consolidation and serialization read.

Parameters:

• distribution (nvalchemi.distributed._core.spec.DistributionSpec) – Required. A DistributionSpec carrying the StoragePolicy and escape-hatch tuples.

• owned_only_outputs (frozenset[str]) – Output keys whose per-atom values are already globally-correct on every rank (e.g. forces computed from replicated global state like Ewald/PME reciprocal S(k)). Consolidation slices these to [:n_owned] rather than halo-reverse-summing.

• all_reduce_outputs (frozenset[str]) – Output keys whose value on each rank is a partial contribution that must be summed across the mesh to give the globally-correct value.

• output_kinds (dict[str, nvalchemi.distributed.output_kinds.OutputKind]) – Per-output classification (OutputKind) consumed by output consolidation. Outputs missing from this dict fall back to a shape-based heuristic (shape[0] == n_padded ⇒ per-atom) and emit a one-shot warning.

• system_reductions (bool)

• node_energy_key (str | None)

• node_virial_key (str | None)

• gp_replicate_geometry (bool)

• outputs (dict[str, OutputSpec] | None)

• compile (CompilePolicy | None)

classmethod from_dict(d)

Inverse of to_dict().

Resolves op qualnames — the caller must ensure the relevant op-registering modules have been imported first.

Parameters:

d (dict[str, Any])

Return type:

MLIPSpec

classmethod load(path)

Load a spec previously saved via save().

Parameters:

path (str | Any)

Return type:

MLIPSpec

merge(other)

Merge two specs for a composed pipeline.

Storage policy: merged via _merge_policies() (two halo policies keep halo and take the more permissive scatter/gather mode). Output-classification sets: union. Escape-hatch tuples: concatenated. system_reductions: logical OR. compile: kept only when both sides declare a compatible CompilePolicy (see _merge_compile_policies()), else None with a warning.

Parameters:

other (MLIPSpec)

Return type:

MLIPSpec

save(path)

Write the spec to path as JSON.

Parameters:

path (str | Any)

Return type:

None

to_dict()

Serialize to a JSON-friendly dict.

Schema:

{
  "version": 2,
  "core": <DistributionSpec.to_dict()>,
  "system_reductions": bool,
  "owned_only_outputs": [...],
  "all_reduce_outputs": [...],
  "output_kinds": [[key, kind], ...]
}

Op handles are encoded as "<namespace>::<name>" strings, resolved at load time provided the registering module has been imported.

Return type:

dict[str, Any]

with_adapters(*adapters)

Return a copy with adapters added to the distribution.

Each adapter is lowered onto custom_ops (OpAdapter) or third_party_helpers (everything else), composing with whatever the spec already declares. Lets a model take a preset and attach model-discovered adapters without rebuilding the spec by hand. All other settings are carried unchanged.

Parameters:

adapters (Any)

Return type:

MLIPSpec

with_compile(policy)

Return a copy with the CompilePolicy set (replacing any existing one). All other settings are carried unchanged.

Parameters:

policy (CompilePolicy)

Return type:

MLIPSpec

with_outputs(outputs)

Return a copy with outputs layered over the declared ones.

Composes additively, like with_adapters(): keys given here replace their counterparts and every other classification is carried through.

Parameters:

outputs (dict[str, OutputSpec]) – Classifications to override.

Returns:

The spec with those outputs reclassified.

Return type:

MLIPSpec