torchcell.models.self_attention_deep_set.SelfAttention

class SelfAttention(dim_in: int, dim_out: int, num_heads: int = 1)

Bases: Module

Multi-head self-attention returning outputs and per-graph weights.

forward(x: Tensor, batch: Tensor) → tuple[Tensor, list[Tensor]]

Attend over node features, returning outputs and per-graph weights.