enerzyme.models.mace.core.MACEWrapper

class enerzyme.models.mace.core.MACEWrapper(atomic_numbers: List[int], r_max: float, num_bessel: int, num_polynomial_cutoff: int, interaction_cls: str, interaction_cls_first: str, max_ell: int, correlation: int, num_interactions: int, MLP_irreps: str, radial_MLP: List[int], hidden_irreps: str, gate: str, avg_num_neighbors: float)

Bases: BaseFFCore

__init__(atomic_numbers: List[int], r_max: float, num_bessel: int, num_polynomial_cutoff: int, interaction_cls: str, interaction_cls_first: str, max_ell: int, correlation: int, num_interactions: int, MLP_irreps: str, radial_MLP: List[int], hidden_irreps: str, gate: str, avg_num_neighbors: float)

build(built_layers) → None

get_output(Ra: Tensor, Za: Tensor, batch_seg: Tensor) → Dict[str, Tensor]