ferritin_plms/esm3/layers/
transformer_stack.rs

1//! ESM3 TransformerStack: stacked UnifiedTransformerBlocks with geometric attention.
2
3use crate::esm3::layers::blocks::UnifiedTransformerBlock;
4use crate::esm3::models::esm3::ESM3Config;
5use crate::esm3::utils::affine3d::Affine3D;
6use candle_core::{Module, Result, Tensor};
7use candle_nn::{self as nn, VarBuilder};
8
9pub struct TransformerStack {
10    blocks: Vec<UnifiedTransformerBlock>,
11    norm: nn::LayerNorm,
12}
13
14impl TransformerStack {
15    pub fn load(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
16        let mut blocks = Vec::with_capacity(config.n_layers);
17        for i in 0..config.n_layers {
18            blocks.push(UnifiedTransformerBlock::load(
19                vb.pp(format!("blocks.{}", i)),
20                config,
21                i,
22            )?);
23        }
24
25        // Final LayerNorm: weight only, no bias.
26        let norm_weight = vb.pp("norm").get((config.d_model,), "weight")?;
27        let norm = nn::LayerNorm::new_no_bias(norm_weight, 1e-5);
28
29        Ok(Self { blocks, norm })
30    }
31
32    /// Forward pass through the full ESM3 transformer stack.
33    ///
34    /// - `x`:           `(B, L, d_model)` input embeddings.
35    /// - `sequence_id`: optional `(B, L)` int — per-protein bin-packing IDs.
36    /// - `affine`:      optional per-residue local frames for geometric attention.
37    /// - `affine_mask`: optional `(B, L)` u8 — 1 where frame is valid.
38    /// - `chain_id`:    optional `(B, L)` int — chain identity per residue.
39    ///
40    /// Returns `(post_norm, pre_norm)` matching the Python API.
41    pub fn forward(
42        &self,
43        x: &Tensor,
44        sequence_id: Option<&Tensor>,
45        affine: Option<&Affine3D>,
46        affine_mask: Option<&Tensor>,
47        chain_id: Option<&Tensor>,
48    ) -> Result<(Tensor, Tensor)> {
49        let mut x = x.clone();
50
51        for block in &self.blocks {
52            x = block.forward(&x, sequence_id, affine, affine_mask, chain_id)?;
53        }
54
55        let post_norm = self.norm.forward(&x)?;
56        Ok((post_norm, x))
57    }
58}
ferritin_plms/esm3/layers/transformer_stack.rs

ferritin_plms/esm3/layers/
transformer_stack.rs
