Structure module, templates and recycling · INTEGRATION

Recycling representations and clean-template state

HardRaw PyTorchAbout 30 min2 public examples plus private tests

AlphaFold2OpenFoldAlphaFold3 ESMFold RFdiffusion

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Prepare inputs for another model pass using available previous predictions. AF2, ESMFold and AF3 add their declared feedback to fresh base features; clean_template exposes a cached backbone as one template. Return only the chosen mode’s fields plus the next completed-pass count and whether another pass remains.

Function signature

recycling_state_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    mode: str = 'af2',
    max_passes: int = 3,
    eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]

Worked example

Fresh base plus previous-pass feedback

For an AF3 example with one batch, one valid residue and two feature channels, use single_base=[10,20], pair_base=[2,4], and both previous feature vectors [1,3] (with their batch/residue/pair axes). Set single and pair normalization scales to [1,1], shifts to [0,0], w_single=2I, w_pair=4I, eps=3, has_previous=[True], iteration=[1], and max_passes=3. Each normalized previous vector is [-0.5,0.5], so feedback is [-1,1] for single and [-2,2] for pair. The outputs are [9,21] and [0,6]; next_iteration=[2] and continue=[True]. AF3 adds no distance term.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`mode`strDefault `'af2'`	Choose "af2", "esmfold", "af3", "clean_template". Each choice uses the matching input subset and computational rule below.
`max_passes`intDefault `3`	Maximum recycling passes.
`eps`floatDefault `1e-05`	Positive value added to final-axis population variance inside each LayerNorm square root. It is not a mask/count denominator or an output offset.
`glycine_index`intDefault `7`	Vocabulary index for glycine.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Implement one explicit state-preparation step, not a hidden call to a full network. mode is 'af2' (default), 'esmfold', 'af3', or 'clean_template'. iteration $[B]$ is an integer number of already completed passes in $[0,\mathrm{max\_passes}]$ , $\mathrm{max\_passes}\ge 1$ integer default $3$ ; has_previous $[B]$ is Boolean.
$\mathrm{active}=\mathrm{iteration}<\mathrm{max\_passes}$ ; $\mathrm{feedback}=\mathrm{active}$ AND has_previous. Return $\mathrm{next\_iteration}=\mathrm{iteration}+\mathrm{active}$ and $\mathrm{continue}=(\mathrm{next\_iteration}<\mathrm{max\_passes})$ .
This only chooses whether another caller-owned pass is needed; no model evaluation, adaptive stopping criterion or stochastic behavior is inside the function.
For $\mathrm{af2}/\mathrm{esmfold}/\mathrm{af3}$ : single_base,previous_single $[B,N,\mathrm{Cs}]$ , pair_base,previous_pair $[B,N,N,\mathrm{Cz}]$ , residue_mask $[B,N]$ True means valid. Zero masked previous features before normalization.
Independently normalize each previous single/pair feature vector over its channels, population variance, $\sqrt{\mathrm{var}+\varepsilon}$ , $\varepsilon>0$ default $10^{-5}$ , with supplied learned scales/shifts.
Compute $\mathrm{ds}=\operatorname{LN}(\mathrm{previous\_single})\mathbin{@}w_{\mathrm{single}}[\mathrm{Cs},\mathrm{Cs}]$ , $\mathrm{dz}=\operatorname{LN}(\mathrm{previous\_pair})\mathbin{@}w_{\mathrm{pair}}[\mathrm{Cz},\mathrm{Cz}]$ . $\mathrm{AF2}/\mathrm{ESMFold}$ can use identity w_single/w_pair; AF3 uses learned projections.
The af2 single denotes the previous TARGET MSA row, not every MSA row. Output $\mathrm{single}=\mathrm{single\_base}+\mathrm{feedback}\cdot \mathrm{ds}$ and $\mathrm{pair}=\mathrm{pair\_base}+\mathrm{feedback}\cdot \mathrm{dz}$ , masked to exact zero on invalid residues/pairs.
Every iteration starts from the supplied base state; do not add the old state twice. $\mathrm{af2}/\mathrm{esmfold}$ also take atom_positions $[B,N,4,3]$ ordered $N$ ,CA, $C$ ,CB and Boolean atom_mask; missing atoms are ignored. af2 selects CA for $\mathrm{residue\_type}=\mathrm{glycine\_index}$ (default $7$ ), CB otherwise.
esmfold reconstructs $\beta=\mathrm{CA}-0.58273431\cdot \operatorname{cross}(\mathrm{CA}-N,C-\mathrm{CA})+0.56802827\cdot (\mathrm{CA}-N)-0.54067466\cdot (C-\mathrm{CA})$ , requiring $N/\mathrm{CA}/C$ , and does not substitute glycine CA. pseudo_beta is zero unless its atoms, residue and feedback are valid.
For each valid $\beta$ pair form $\mathrm{bin}=\sum(\mathrm{distance\_squared} > \mathrm{distance\_edges}^{2})$ , equality in the lower bin. distance_edges $[K]$ is fixed, positive, strictly increasing, $K\ge 1$ . Add $w_{\mathrm{distance}}[\mathrm{bin}]+b_{\mathrm{distance}}$ to dz only on valid $\beta$ pairs.
Return distance_bins $[B,N,N]$ with $-1$ for invalid/no-feedback pairs. af3 has no geometry tensors/outputs and adds no distance term.
clean_template instead takes only previous_clean $[B,N,3,3] N/\mathrm{CA}/C$ , has_previous and iteration: return template $[B,1,N,3,3]$ equal to the previous clean backbone when feedback is True, otherwise zero, plus the two iteration outputs.
This covers deterministic clean-coordinate cache selection; template_pair_features computes the nontrivial next-step template geometry. No detach is allowed or needed: training stop-gradient boundaries belong to the caller, while this mathematical map stays differentiable.
All floating inputs must be finite; all $B$ , $N$ ,channel dimensions are positive.

Why this matters

Recycling feeds information from an earlier prediction into fresh base representations; it does not itself rerun the network.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def recycling_state_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    mode: str = 'af2',
    max_passes: int = 3,
    eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    single: (B, N, Cs). Returned in state modes (af2, esmfold, af3).
    pair: (B, N, N, Cz). Returned in state modes (af2, esmfold, af3).
    pseudo_beta: (B, N, 3). Returned only for mode="af2" or "esmfold".
    distance_bins: (B, N, N). Returned only for mode="af2" or "esmfold".
    template: (B, 1, N, 3, 3). Returned only for mode="clean_template".
    next_iteration: (B,).
    continue: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement recycling_state_update")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Related exercises

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.