torchcell.models.hetero_cell_bipartite_dango_gi_lazy.GeneInteractionPredictor

class GeneInteractionPredictor(hidden_dim: int, num_heads: int = 4, num_layers: int = 2, dropout: float = 0.1, activation: str = 'relu')

Bases: Module

Predict gene interaction scores from gene embeddings via a HyperSAGNN.

forward(gene_embeddings: Tensor, batch: Tensor | None = None) → Tensor

Score gene interactions from static/dynamic embedding differences.

Parameters:

• gene_embeddings – Tensor of shape [total_genes, hidden_dim]

• batch – Optional tensor [total_genes] indicating batch assignment

Returns:

Tensor of shape [num_batches]

Return type:

interaction_scores