torchcell.models.cell_graph_transformer_metabolism.ProductScalarHead

class ProductScalarHead(hidden_dim: int, use_gene_pool: bool = True, use_pert_pool: bool = True, dropout: float = 0.1, param_dim: int = 1)

Bases: Module

Whole-cell readout for a SCALAR production phenotype (one value per strain).

Reads the CLS token, optionally concatenated with a mean pool over the equivariant perturbed gene embeddings, and emits exactly one value. Structurally this is the S1 pooled readout; it is a distinct class from the vector head so that “this head is scalar” is a property of the module rather than of a config string that some other code has to agree with. output_dim is fixed at 1 by construction – a scalar target assigned into a wider head broadcasts silently, and the cheapest place to make that impossible is here.

forward(h_CLS: Tensor, H_genes_pert: Tensor, pert_pool: Tensor | None = None) → Tensor

Return [batch, 1] (point) or [batch, 1, param_dim] (distributional).