xnn.gnn.models.aimnet2_foundation.load_artifact

xnn.gnn.models.aimnet2_foundation.load_artifact(source, member=0)

Read an AIMNet2 artifact into one mapping.

Parameters:

• source (str or Path) – A reference .pt artifact, a config.json + ensemble_<k>.safetensors directory, or one .safetensors file (its config.json is read from the same directory).

• member (int, optional) – Ensemble member to take from a directory, by default 0.

Returns:

state_dict (CPU tensors), model_yaml and the metadata keys of the reference format (cutoff, implemented_species, coulomb_mode, d3_params, …).

Return type:

dict

Raises:

• FileNotFoundError – If the file, directory or ensemble member does not exist.

• ValueError – If the file is not an AIMNet2 artifact.