Struct AtomEncoder 

pub struct AtomEncoder { /* private fields */ }

3-block atom-level transformer encoder.

Maps per-atom features ([B, N_atom, d_atom=128]) to per-token atom representations ([B, N_tok, d_token=768]) via windowed attention and token-level pooling.

Implementations

impl AtomEncoder

pub fn load( vb: VarBuilder<'_>, d_atom: usize, d_token: usize, n_blocks: usize, ) -> Result<Self>

Load the atom encoder from a VarBuilder rooted at inputs.atom_encoder.*.

Arguments

• vb — builder rooted at inputs.atom_encoder
• d_atom — per-atom feature dimension (128)
• d_token — per-token output dimension (768)
• n_blocks — number of atom transformer blocks (3)

pub fn forward(&self, x: &Tensor) -> Result<Tensor>

Encodes per-atom features to per-token atom representations.

• Input: [B, N_atom, d_atom=128]
• Output: [B, N_tok, d_token=768]

The number of output tokens depends on the residue-level grouping of atoms.