torchcell.models.hetero_cell_bipartite_dango_gi.GeneInteractionPredictor
- class GeneInteractionPredictor(hidden_dim: int, num_heads: int = 4, num_layers: int = 2, dropout: float = 0.1)[source]
Bases:
ModulePredict gene-interaction scores from squared static/dynamic embedding diffs.
- forward(gene_embeddings: Tensor, batch: Tensor | None = None) Tensor[source]
Return per-batch gene-interaction scores.
- 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