torchcell.nn.nsa_encoder.NSAEncoder

class NSAEncoder(input_dim: int, hidden_dim: int, pattern: list[Literal['M', 'S']] | None = None, num_heads: int = 8, dropout: float = 0.1, activation: Module = GELU(approximate='none'))

Bases: Module

NSA encoder applying an input projection then masked/self-attention blocks.

Applies an input projection followed by a sequence of masked (MAB) and self-attention (SAB) blocks.

forward(x: Tensor, data_or_edge_index: Tensor | HeteroData, edge_attr: Tensor | None = None, batch: Tensor | None = None) → Tensor

Encode node features, padding per graph and unpadding the output.