torchcell.nn
Neural-network layers shared by the models. hetero_nsa, nsa_encoder, masked_attention_block and self_attention_block build node-set attention over graphs with FlexAttention adjacency masks; masked_gin_conv is a GIN convolution for the lazy, masked subgraph representation; and stoichiometric_hypergraph_conv is a stoichiometry-aware hypergraph convolution over metabolic networks.
hetero_nsa
Heterogeneous Node-Set Attention blocks and encoder over HeteroData graphs.
Classes
Heterogeneous Node-Set Attention module stacking per-pattern blocks. |
|
Full HeteroNSA encoder with input projections and stacked layers. |
masked_attention_block
Masked attention blocks using FlexAttention with adjacency and edge masks.
Classes
Memory-efficient Masked Attention Block that uses FlexAttention on GPU. |
|
Masked attention block with learned edge attribute projections. |
masked_gin_conv
Masked GIN convolution for lazy subgraph representation.
Classes
Graph Isomorphism Network with masked message passing. |
nsa_encoder
Node self-attention encoder stacking masked and self-attention blocks.
Classes
NSA encoder applying an input projection then masked/self-attention blocks. |
self_attention_block
Device-adaptive self-attention block used across torchcell models.
Classes
Self-Attention Block (SAB) with device-adaptive implementation. |
stoichiometric_hypergraph_conv
Stoichiometry-aware hypergraph convolution layer for metabolic networks.
Classes
Hypergraph convolution with optional stoichiometric gating and attention. |