torchcell.models.hetero_cell_bipartite_dango_gi.SelfAttentionGraphAggregation

class SelfAttentionGraphAggregation(hidden_dim: int, num_graphs: int, num_heads: int = 4, dropout: float = 0.0)

Bases: Module

Self-attention aggregation over multiple graph views of the same nodes.

forward(graph_outputs: dict[str, Tensor]) → tuple[Tensor | None, Tensor | None]

Self-attend across graph views and mean-pool to aggregated node features.

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