Struct TransformerStack 

pub struct TransformerStack { /* private fields */ }

Implementations

impl TransformerStack

pub fn load(vb: VarBuilder<'_>, config: &ESM3Config) -> Result<Self>

pub fn load_geometric_encoder( vb: VarBuilder<'_>, config: &ESM3Config, ) -> Result<Self>

The stack used by the ESM3 structure encoder.

Two differences from load, both read off the released checkpoint rather than assumed: every block is geometric-only (there are no attn.* tensors), and there is no trailing norm.weight — the encoder’s final normalization is the identity.

pub fn forward( &self, x: &Tensor, sequence_id: Option<&Tensor>, affine: Option<&Affine3D>, affine_mask: Option<&Tensor>, chain_id: Option<&Tensor>, ) -> Result<(Tensor, Tensor)>

Forward pass through the full ESM3 transformer stack.

• x: (B, L, d_model) input embeddings.
• sequence_id: optional (B, L) int — per-protein bin-packing IDs.
• affine: optional per-residue local frames for geometric attention.
• affine_mask: optional (B, L) u8 — 1 where frame is valid.
• chain_id: optional (B, L) int — chain identity per residue.

Returns (post_norm, pre_norm) matching the Python API.