torchcell.models.hetero_cell_bipartite_dango_gi.DangoLikeHyperSAGNN

class DangoLikeHyperSAGNN(hidden_dim: int, num_heads: int = 4, num_layers: int = 2, dropout: float = 0.1)

Bases: Module

Dango-like HyperSAGNN for local gene interaction prediction.

Implements multi-layer multi-head self-attention with ReZero connections.

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

Return static and attention-refined dynamic gene 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] dynamic_embeddings: Tensor of shape [total_genes, hidden_dim]

Return type:

static_embeddings