xnn.common.models.hub.checkpoint.build_potential#

xnn.common.models.hub.checkpoint.build_potential(config, state_dict, *, dtype=None, dispersion=None, compute_forces=True, compute_stress=False, eeq_reuse=False, model_options=None, head=None, label='checkpoint')[source]#

Rebuild a model from its config and weights.

The model is built in float64, so the constant tables of the physics terms (D3 / D4 reference data, LES kernels) hold their exact values, and cast once afterwards, to dtype or to the weights’ own dtype. Building in float32 and upcasting would keep float32-rounded tables, which costs about 5e-8 hartree in a D4 energy even when serving in float64.

A config that records subtracted_dispersion (what the training labels had removed) gets that term added back unless dispersion=False; see resolve_dispersion().

A fine-tuned checkpoint is served as a plain potential: of a multi-head model (extra["heads"]) the head named by head is kept (select()), and LoRA updates (extra["lora"]) are folded into the base weights (merge_lora()).

Parameters:

  • config (Config) – The run configuration.

  • state_dict (dict) – Weights in the ForceStressOutput layout.

  • dtype (torch.dtype or str, optional) – Dtype to serve in. Defaults to the dtype of the weights.

  • dispersion (dict, str, bool or None, optional) – D3 / D4 correction to add (as in extra["dispersion"]); None adds the recorded one, if any, and False adds nothing.

  • compute_forces (bool, optional) – Heads of the returned ForceStressOutput.

  • compute_stress (bool, optional) – Heads of the returned ForceStressOutput.

  • eeq_reuse (bool, optional) – Carry the D4 EEQ solve over between steps (molecular dynamics).

  • model_options (dict, optional) – Overrides of config.model.extra applied before the model is built: deployment knobs that change no weights, such as AIMNet2’s Coulomb method ({"coulomb": "dsf", "lr_cutoff": 15.0}); the weights still load strictly.

  • head (str, optional) – The head of a multi-head checkpoint to serve; required when it has several, ignored (if it matches) for a single-head one.

  • label (str, optional) – Name used in log messages.

Returns:

The model with its weights loaded, in eval mode.

Return type:

ForceStressOutput

Raises:

ValueError – If dispersion is given for a model that already includes one, or head is missing or unknown for a multi-head checkpoint.

On this page
Show Source