torchcell.models.hetero_cell_bipartite_dango.GeneInteractionAttention

class GeneInteractionAttention(hidden_dim: int, num_heads: int = 8, dropout: float = 0.1)

Bases: Module

Self-attention for digenic and trigenic gene interaction prediction.

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

Apply masked self-attention per batch to produce dynamic embeddings.

Parameters:

• gene_embeddings – Tensor of shape [total_genes, hidden_dim].

• batch – Optional tensor [total_genes] indicating batch assignment.

Returns:

Tensor of shape [total_genes, hidden_dim].

Return type:

dynamic_embeddings