torchcell.models.hetero_cell_bipartite_dango_gi_lazy.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)[source]

Bases: Module

HeteroConv wrapper with configurable aggregation strategies.

Supports: sum, mean, cross_attention, pairwise_interaction

forward(x_dict: dict[str, Tensor], edge_index_dict: dict[Tuple[str, str, str], Tensor], edge_mask_dict: dict[Tuple[str, str, str], Tensor] | None = None) → tuple[dict[str, Tensor], dict[str, Tensor | None] | None][source]

Apply graph convolutions and aggregate using specified method.

LAZY VERSION: Accepts edge_mask_dict for masking edges during message passing.

Parameters:

  • x_dict – Node features by node type

  • edge_index_dict – Edge indices by edge type

  • edge_mask_dict – Optional edge masks by edge type (True = keep edge)

Returns:

Updated node features attention_weights: Graph aggregation weights (if applicable)

Return type:

out_dict