MolecularDiffusion.modules.models.apo2mol.attn¶
Retrieval-augmented linear attention over the ligand token block.
Ported from others/Apo2Mol/models/attn.py:17-44.
The retrieval path itself is out of scope (topk_prompt: 0 in the
released configs/training.yaml, so
prompt_hbap_ligand_batch_all_list is always empty and the module is
called with h_retrieved = h, i.e. self-attention). The module is still
ported because its four weight matrices are in the released checkpoint and
they do affect the forward pass in that degenerate configuration.
CrossAttention from the same upstream file is not ported: it is imported
by models/molopt_score_model.py:15 but never instantiated, so it
contributes no checkpoint tensors and has no call site.
Classes¶
Linear attention mixing a token block with a retrieved block. |
Module Contents¶
- class MolecularDiffusion.modules.models.apo2mol.attn.RetAugmentationLinearAttention(in_dim: int, d: int, context_dim: int)¶
Bases:
torch.nn.ModuleLinear attention mixing a token block with a retrieved block.
handh_retrievedare(B, N, in_dim)– note this is the one place in the model that uses a padded ligand block rather than the flat scatter layout (molopt_score_model.py:390-414pads to 150 atoms and unpads afterwards).- forward(h: torch.Tensor, h_retrieved: torch.Tensor) torch.Tensor¶
- cond_flag = False¶
- linear_attn¶
- out¶
- to_k¶
- to_v¶