torchcell.models.dango.HyperSAGNN

class HyperSAGNN(hidden_channels: int, num_heads: int = 4)

Bases: Module

Fully vectorized Hypergraph Self-Attention Graph Neural Network that handles all perturbation sets in a single forward pass.

forward(embeddings: Tensor, batch_indices: Tensor, num_sets: int) → Tensor

Forward pass processing all nodes at once with masked attention.

Parameters:

• embeddings – Tensor of shape [total_nodes, hidden_channels]

• batch_indices – Tensor of shape [total_nodes] indicating set membership, with ids in [0, num_sets)

• num_sets – Number of sets (genotypes) in the batch; the output has one score per set, ordered by set id.

Returns:

Predicted interaction scores with shape [num_sets]

Raises:

ValueError – if a set id is outside [0, num_sets), or a set has no node. A set with no perturbed gene has nothing to attend over and no interaction score, so it is refused rather than scored 0.