torchcell.models.hetero_cell_nsa_retry.AttentionalGraphAggregation

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

Bases: Module

Attentional aggregation for graph-level representations.

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

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