xnn.common.models.registry#

Name -> class registry so new models are added without touching the core.

For example:

from xnn.common.models.registry import register_model, build_model

@register_model("mymodel")
class MyModel(InteratomicPotential):
    @classmethod
    def from_config(cls, cfg): ...

A model becomes usable from any config (YAML/Hydra/argparse) the moment it is imported and registered. This is the extension point referenced in the README.

Module Attributes

FINETUNE_KEYS

model.extra keys that describe how a pretrained model is adapted rather than its architecture.

Functions

add_dispersion(model, spec)

Wrap model in a D3 or D4 dispersion correction.

available_models()

List the names of all currently registered models.

build_model(cfg)

Instantiate a registered model from a configuration dataclass.

core_model(model)

The innermost potential under the force, dispersion, LES and multi-head wrappers.

deployment_config(cfg)

A copy of cfg describing the deployed (single-head, LoRA-merged) model.

has_electrostatics(model)

Whether a model adds the Coulomb energy of charges it predicts itself.

prepare_model(cfg)

Build a model and the config that rebuilds it.

recorded_dispersion(cfg)

The subtracted_dispersion record of a (possibly older) Config.

register_model(name)

Return a class decorator that registers a model under name.

replace_module(chain, old, new)

Replace old by new inside a chain of wrappers linked by .model.

resolve_dispersion(explicit, recorded)

The dispersion a deployment adds to a checkpoint's model.
Show Source