torchcell.models.equivariant_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. Any integer ids; they need not be contiguous or start at 0.

Returns:

Perturbation representations with shape [num_batches, hidden_channels], one row per DISTINCT set id, in ascending id order.