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