torchcell.models.hetero_cell_bipartite_dango_gi_lazy.SelfAttentionGraphAggregation
- class SelfAttentionGraphAggregation(hidden_dim: int, num_graphs: int, num_heads: int = 4, dropout: float = 0.0)[source]
Bases:
ModuleSelf-attention mechanism for aggregating multiple graph representations of the same nodes.
- forward(graph_outputs: dict[str, Tensor]) tuple[Tensor | None, Tensor | None][source]
Attend across per-graph node features and mean-pool to one representation.
- Parameters:
graph_outputs – Dict mapping graph names to node features [num_nodes, hidden_dim]
- Returns:
Aggregated node features [num_nodes, hidden_dim] attention_weights: Attention weights [num_nodes, num_graphs, num_graphs]
- Return type:
aggregated