xnn.gnn.models.mace_foundation.from_mace_torch#

xnn.gnn.models.mace_foundation.from_mace_torch(upstream, head=None, dtype=None)[source]#

Convert a mace-torch model (plain or ScaleShift) to an xnn MACE.

Reads every architecture hyper-parameter off the upstream module, builds the equivalent MACE, and copies the weights; multi-head checkpoints are sliced to the requested head. The converted model reproduces the upstream energies, forces and stress to numerical precision (see tests/test_mace.py and the foundation fidelity notebook).

Parameters:

upstream (torch.nn.Module) – A mace.modules.models.MACE or ScaleShiftMACE instance.
head (str, optional) – Head to keep for multi-head checkpoints (see _resolve_head()).
dtype (torch.dtype or str, optional) – Final dtype; None keeps the upstream parameters’ dtype.

Returns:

The converted model.

Return type:

MACE

Raises:

NotImplementedError – For upstream features outside the xnn MACE (un-enveloped radial embeddings, unknown interaction blocks / radial bases / transforms, readout layouts other than one readout per interaction).