torchcell.models.equivariant_cell_graph_transformer.PerGeneHead

class PerGeneHead(hidden_dim: int, output_dim: int = 1, dropout: float = 0.1, param_dim: int = 1, free_gene_dim: int = 0, num_genes: int = 0, in_mult: int = 1, extra_dim: int = 0, film_dim: int = 0)

Bases: Module

Per-gene (equivariant) readout head for gene-resolved phenotypes.

Maps the equivariant perturbed gene embeddings H_genes_pert [batch, N, d] to a per-gene prediction, preserving the per-node structure. Used for expression (MicroarrayExpressionPhenotype / RnaseqExpressionPhenotype) AND proteome (ProteinAbundancePhenotype).

graph_level == “node” selects this head.

forward(H_genes_pert: Tensor, film_cond: Tensor | None = None) → Tensor

Forward pass of the per-gene head.

Parameters:

• H_genes_pert – [batch, N, d] equivariant perturbed gene embeddings.

• film_cond – [batch, film_dim] strain-level conditioning vector. Required for FiLM to engage: the scale/shift are only applied when the head was built with film_dim > 0 AND this is not None, so passing None reduces the head to its unconditioned form.

Returns:

[batch, N] if output_dim == param_dim == 1; [batch, N, param_dim] for a distributional scalar-per-gene head; else [batch, N, output_dim (, param_dim)].

Return type:

predictions