torchcell.models.cell_graph_transformer.HyperSAGNN

class HyperSAGNN(hidden_channels: int, num_heads: int = 4)[source]

Bases: Module

Hypergraph Self-Attention Graph Neural Network for perturbation sets.

Adapted from DANGO model to compute perturbation representations via masked self-attention within perturbation sets.

forward(embeddings: Tensor, batch_indices: Tensor) → Tensor[source]

Forward pass processing perturbed genes with masked attention.

Parameters:

  • embeddings – Tensor of shape [total_pert_genes, hidden_channels]

  • batch_indices – Tensor of shape [total_pert_genes] indicating set membership

Returns:

Perturbation representations with shape [num_batches, hidden_channels]