enerzyme.models.layers.electron_embedding.BaseElectronEmbedding#

class enerzyme.models.layers.electron_embedding.BaseElectronEmbedding(dim_embedding: int, num_residual: int, attribute: Literal['charge', 'spin'] = 'charge')[source]#

Bases: BaseFFLayer

__init__(dim_embedding: int, num_residual: int, attribute: Literal['charge', 'spin'] = 'charge') → None[source]#
abstractmethod get_electron_embedding(atom_embedding, Q, batch_seg) → Dict[Literal['electron_embedding'], Tensor][source]#