torchcell.models.cell_graph_transformer.GraphRegularizedTransformerLayer

class GraphRegularizedTransformerLayer(hidden_dim: int, num_heads: int, adjacency_matrices: dict[str, Tensor] | None = None, regularized_head_config: dict[str, dict[str, Any]] | None = None, dropout: float = 0.1)

Bases: Module

Transformer layer with graph-regularized attention heads.

Uses manual attention computation to get both output and attention weights for graph regularization loss.

forward(x: Tensor, return_attention: bool = False) → tuple[Tensor, Tensor | None]

Forward pass with manual attention computation.

Parameters:

• x – [batch, N+1, d] where N+1 includes CLS token at position 0

• return_attention – Whether to return attention weights for regularization

Returns:

[batch, N+1, d] transformed features gene_attention: [batch, heads, N, N] gene-gene attention weights (if return_attention=True)

Return type:

output