torchcell.models.hetero_cell_bipartite_dango_gi_lazy.DangoLikeHyperSAGNN

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

Bases: Module

Dango-like HyperSAGNN for local gene interaction prediction.

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

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

Compute 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