torchcell.models.hetero_cell_bipartite_dango_gi.AttentionalGraphAggregation

class AttentionalGraphAggregation(in_channels: int, out_channels: int, dropout: float = 0.1)

Bases: Module

Gated attentional pooling of node features into a graph-level representation.

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

Attentionally aggregate node features x grouped by index.