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    /// `None` for the structure encoder, which ends without a final norm.
12    norm: Option<nn::LayerNorm>,
13}
14
15impl TransformerStack {
16    pub fn load(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
17        let mut blocks = Vec::with_capacity(config.n_layers);
18        for i in 0..config.n_layers {
19            blocks.push(UnifiedTransformerBlock::load(
20                vb.pp(format!("blocks.{}", i)),
21                config,
22                i,
23            )?);
24        }
25
26        // Final LayerNorm: weight only, no bias.
27        let norm_weight = vb.pp("norm").get((config.d_model,), "weight")?;
28        let norm = Some(nn::LayerNorm::new_no_bias(norm_weight, 1e-5));
29
30        Ok(Self { blocks, norm })
31    }
32
33    /// The stack used by the ESM3 structure encoder.
34    ///
35    /// Two differences from [`load`][Self::load], both read off the released
36    /// checkpoint rather than assumed: every block is geometric-only (there
37    /// are no `attn.*` tensors), and there is no trailing `norm.weight` —
38    /// the encoder's final normalization is the identity.
39    pub fn load_geometric_encoder(vb: VarBuilder, config: &ESM3Config) -> Result<Self> {
40        let mut blocks = Vec::with_capacity(config.n_layers);
41        for i in 0..config.n_layers {
42            blocks.push(UnifiedTransformerBlock::load_geometric(
43                vb.pp(format!("blocks.{}", i)),
44                config,
45            )?);
46        }
47        Ok(Self { blocks, norm: None })
48    }
49
50    /// Forward pass through the full ESM3 transformer stack.
51    ///
52    /// - `x`:           `(B, L, d_model)` input embeddings.
53    /// - `sequence_id`: optional `(B, L)` int — per-protein bin-packing IDs.
54    /// - `affine`:      optional per-residue local frames for geometric attention.
55    /// - `affine_mask`: optional `(B, L)` u8 — 1 where frame is valid.
56    /// - `chain_id`:    optional `(B, L)` int — chain identity per residue.
57    ///
58    /// Returns `(post_norm, pre_norm)` matching the Python API.
59    pub fn forward(
60        &self,
61        x: &Tensor,
62        sequence_id: Option<&Tensor>,
63        affine: Option<&Affine3D>,
64        affine_mask: Option<&Tensor>,
65        chain_id: Option<&Tensor>,
66    ) -> Result<(Tensor, Tensor)> {
67        let mut x = x.clone();
68
69        for block in &self.blocks {
70            x = block.forward(&x, sequence_id, affine, affine_mask, chain_id)?;
71        }
72
73        // The structure encoder has no final norm, so this is the identity.
74        let post_norm = match &self.norm {
75            Some(norm) => norm.forward(&x)?,
76            None => x.clone(),
77        };
78        Ok((post_norm, x))
79    }
80}
ferritin_plms/esm3/layers/transformer_stack.rs

ferritin_plms/esm3/layers/
transformer_stack.rs
