torchcell.models.hetero_cell_bipartite_dango_gi.HeteroConvAggregator
- class HeteroConvAggregator(convs: dict[Tuple[str, str, str], Module], hidden_channels: int, aggregation_method: str = 'cross_attention', aggregation_config: dict[str, Any] | None = None, aggregation_norm: str | None = None)[source]
Bases:
ModuleHeteroConv wrapper with configurable aggregation strategies.
Supports sum, mean, cross_attention, and pairwise_interaction.
- forward(x_dict: dict[str, Tensor], edge_index_dict: dict[Tuple[str, str, str], Tensor]) tuple[dict[str, Tensor], dict[str, Tensor | None] | None][source]
Apply per-edge-type convolutions and aggregate using the chosen method.
- Returns:
Updated node features attention_weights: Graph aggregation weights (if applicable)
- Return type:
out_dict