torchcell.nn.hetero_nsa.HeteroNSAEncoder

class HeteroNSAEncoder(input_dims: dict[str, int], hidden_dim: int, node_types: set[str], edge_types: set[tuple[str, str, str]], pattern: list[str], num_layers: int = 3, num_heads: int = 8, dropout: float = 0.1, activation: Module = GELU(approximate='none'), aggregation: Literal['sum', 'mean', 'attention'] = 'sum')

Bases: Module

Full HeteroNSA encoder with input projections and stacked layers.

The constructor accepts a ‘pattern’ (a list of ‘M’/’S’ strings) and a ‘num_layers’ count; a separate HeteroNSA is built for each layer.

forward(data: HeteroData) → tuple[dict[str, Tensor], Tensor]

Project, encode, and pool node embeddings into per-type and graph outputs.