torchcell.models.hetero_cell_bipartite_dango_gi_lazy.AttentionalGraphAggregation

class AttentionalGraphAggregation(in_channels: int, out_channels: int, dropout: float = 0.1, activation: str = 'relu')

Bases: Module

Attentional aggregation pooling node features into graph-level vectors.

forward(x: Tensor, index: Tensor, dim_size: int | None = None) → Tensor

Aggregate node features into graph-level vectors using attention weights.