enerzyme.models.xpainn.core.XPaiNNWrapper

class enerzyme.models.xpainn.core.XPaiNNWrapper(cutoff_sr: float, embed_basis: str, aux_basis: str, node_dim: int, edge_irreps: str, hidden_dim: int, hidden_irreps: str, rbf_kernel: str, num_basis: int, cutoff_fn: str, max_edges: int, action_blocks: int, activation: str, norm_type: str, output_mode: str, output_dim: int, node_average: bool | float, default_dtype: str)

Bases: BaseFFCore

__init__(cutoff_sr: float, embed_basis: str, aux_basis: str, node_dim: int, edge_irreps: str, hidden_dim: int, hidden_irreps: str, rbf_kernel: str, num_basis: int, cutoff_fn: str, max_edges: int, action_blocks: int, activation: str, norm_type: str, output_mode: str, output_dim: int, node_average: bool | float, default_dtype: str)

build(built_layers: List[Module])

get_output(Ra, Za, batch_seg, idx_i, idx_j)