torchcell.nn.hetero_nsa.HeteroNSA

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

Bases: Module

Heterogeneous Node-Set Attention module stacking per-pattern blocks.

Instead of reusing a single set of modules for every block, this version builds a stack of independent blocks, one per element in the provided pattern.

forward(node_embeddings: dict[str, Tensor], data: HeteroData, batch_idx: dict[str, Tensor] | None = None) → dict[str, Tensor]

Apply each block in sequence to the cloned node embeddings.