enerzyme.models.mace.core.MACECore

class enerzyme.models.mace.core.MACECore(max_Za: int, max_ell: int, dim_embedding: int, num_rbf: int, additional_hidden_irreps: str, interaction_cls_first: Literal['RealAgnosticResidualInteractionBlock'], interaction_cls: Literal['RealAgnosticResidualInteractionBlock'], correlation: int | List[int], num_interactions: int, avg_num_neighbors: float, MLP_irreps: str, radial_MLP: List[int], gate: str, shallow_ensemble_size: int = 1)

Bases: BaseFFCore

__init__(max_Za: int, max_ell: int, dim_embedding: int, num_rbf: int, additional_hidden_irreps: str, interaction_cls_first: Literal['RealAgnosticResidualInteractionBlock'], interaction_cls: Literal['RealAgnosticResidualInteractionBlock'], correlation: int | List[int], num_interactions: int, avg_num_neighbors: float, MLP_irreps: str, radial_MLP: List[int], gate: str, shallow_ensemble_size: int = 1)

build(built_layers) → None

get_output(Za: Tensor, vij_sr: Tensor, idx_i_sr: Tensor, idx_j_sr: Tensor, rbf: Tensor, atom_embedding: Tensor, charge_embedding: Tensor | None = None, spin_embedding: Tensor | None = None) → Dict[str, Tensor]