Struct GeometricReasoningOriginalImpl 

pub struct GeometricReasoningOriginalImpl { /* private fields */ }

Implementations

impl GeometricReasoningOriginalImpl

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

Load geometric attention from a VarBuilder rooted at geom_attn.

Weight layout, verified against both esm3_sm_open_v1.pth and esm3_structure_encoder_v0.pth (ferritin-100.21):

s_norm.weight[/.bias]      — (d_model,)
proj.weight[/.bias]        — (15 * v_heads, d_model)
out_proj.weight[/.bias]    — (d_model, 3 * v_heads)
distance_scale_per_head    — (v_heads,)
rotation_scale_per_head    — (v_heads,)

The main ESM3 model stores these without bias terms while the VQ-VAE structure encoder stores them with bias, so config.bias selects.

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

Geometric attention forward pass.

• s: (B, L, d_model) hidden states.
• affine: per-residue local frames (B, L, 3, 3) rot + (B, L, 3) trans.
• affine_mask: (B, L) u8 — 1 where the frame is valid, 0 for frameless positions.
• sequence_id: optional (B, L) int — positions with the same ID form one protein.
• chain_id: optional (B, L) int — positions in the same chain.

Returns (B, L, d_model).