Struct DiffusionModule 

pub struct DiffusionModule { /* private fields */ }

AF3-style EDM diffusion module.

Denoises atom coordinates conditioned on single and pair representations produced by the folding trunk. Uses 12 token-level transformer blocks with pair bias conditioning, followed by coordinate output projection.

Implementations

impl DiffusionModule

pub fn load( vb: VarBuilder<'_>, c_token: usize, c_atom: usize, d_single: usize, d_pair: usize, n_token_blocks: usize, n_token_heads: usize, fourier_dim: usize, ) -> Result<Self>

Load the diffusion module from a VarBuilder rooted at structure_head.*.

Arguments

• vb — builder rooted at structure_head
• c_token — token channel dimension (768)
• c_atom — atom channel dimension (128, reserved for atom transformer TODO)
• d_single — single repr dimension from folding trunk (384)
• d_pair — pair repr dimension from folding trunk (256)
• n_token_blocks — number of token transformer blocks (12)
• n_token_heads — number of heads per token block (16)
• fourier_dim — Fourier noise embedding dimension (256)

pub fn forward( &self, single: &Tensor, pair: &Tensor, n_atoms: usize, num_steps: usize, ) -> Result<Tensor>

Run the AF3-style EDM denoising loop.

Arguments

• single — single representation [B, N_tok, d_single]
• pair — pair representation [B, N_tok, N_tok, d_pair]
• n_atoms — number of atoms (= N_tok for Cα-only, otherwise N_tok * atoms_per_tok)
• num_steps — number of denoising steps (14 fast / 50 quality)

Returns

Token-level coordinates [B, N_tok, 3] in Ångströms (Cα output).