Structure module, templates and recycling

AlphaFold 3 Template Pair Features

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

Describe each ordered template residue pair by distance, residue identities and the direction to its partner in its local backbone frame. Project that geometry, then add the normalized current query-pair projection to every template embedding. Return both the raw geometry features and the query-conditioned embeddings.

Function signature

af3_template_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
    type_count: int = 22,
    norm_eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]

Worked example

The first residue uses CB=[0,0,1]; glycine uses CA=[3,0,0]. Their squared distance is 10, above the squared edge 4. Both local frames align with the global axes. Illustrative norm_eps=3 makes every query projection equal 1. The three valid embeddings are 2.5, 3.5+3/sqrt(9.01), and 3.5-3/sqrt(9.01); pair (1,1) is masked to zero after all biases. Raw features have shape (1,1,2,2,11).

import torch

activations = {
    'distance_edges': torch.tensor([2.0], dtype=torch.float64, requires_grad=True),
    'query': torch.tensor([[[[1.0, 3.0], [1.0, 3.0]], [[1.0, 3.0], [1.0, 3.0]]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'w_features': torch.tensor([[1.0], [2.0], [0.0], [0.0], [0.0], [0.0], [0.0], [1.0], [0.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'b_features': torch.tensor([0.5], dtype=torch.float64, requires_grad=True),
    'query_scale': torch.tensor([1.0, 1.0], dtype=torch.float64, requires_grad=True),
    'query_bias': torch.tensor([0.0, 0.0], dtype=torch.float64, requires_grad=True),
    'w_query': torch.tensor([[1.0], [3.0]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'atom_mask': torch.full((1, 1, 2, 4), True, dtype=torch.bool),
    'pair_mask': torch.tensor([[[True, True], [True, False]]], dtype=torch.bool),
}
structure = {
    'atom_positions': torch.tensor([[[[[0.0, 2.0, 0.0], [0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 0.0, 1.0]], [[3.0, 2.0, 0.0], [3.0, 0.0, 0.0], [5.0, 0.0, 0.0], [3.0, 0.0, 1.0]]]]], dtype=torch.float64, requires_grad=True),
    'residue_type': torch.tensor([[[0, 1]]], dtype=torch.int64),
}
eps = 0.01
type_count = 2
norm_eps = 3.0
glycine_index = 1

result = af3_template_pair_features(activations=activations, parameters=parameters, masks=masks, structure=structure, eps=eps, type_count=type_count, norm_eps=norm_eps, glycine_index=glycine_index)

Expected return

{
    'pseudo_beta': torch.tensor([[[[0.0, 0.0, 1.0], [3.0, 0.0, 0.0]]]], dtype=torch.float64),
    'pair_mask': torch.tensor([[[[True, True], [True, False]]]], dtype=torch.bool),
    'features': torch.tensor([[[[[1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0], [0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.9994449069791544, 0.0, 0.0, 1.0]], [[0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, -0.9994449069791544, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]]]], dtype=torch.float64),
    'embedding': torch.tensor([[[[[2.5], [4.4994449069791544]], [[2.5005550930208456], [0.0]]]]], dtype=torch.float64),
}

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.
`eps`floatDefault `1e-08`	Positive geometric squared-norm floor for frame/angle construction and an additive squared-distance stabilizer for local unit directions; use the distinct placements in the selected feature rule.
`type_count`intDefault `22`	Size of the residue-type vocabulary.
`norm_eps`floatDefault `1e-05`	Positive value added to final-axis population variance inside the query LayerNorm square root; used by the AF3 query-conditioning projection.
`glycine_index`intDefault `7`	Vocabulary index for glycine.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

atom_positions is $[B,T,N,4,3]$ , with atom slots N, CA, C, CB. atom_mask marks available atoms. residue_type uses type_count categories (default 22); one-hot feature order follows these supplied integer IDs.
Pseudo-beta uses CA for glycine(residue_type == glycine_index, default 7) and supplied CB otherwise, with that selected atom’s mask.
A candidate pair requires both complete N/CA/C backbones, both valid pseudo-beta atoms, and pair_mask. Squared-distance bin index counts distance_edges² strictly below squared pseudo-beta distance.
Build the frame at residue i with x along C−CA, y along the component of N−CA orthogonal to x, and z=x×y. Unit-vector normalization divides by sqrt(max(norm²,eps)); both endpoint frames require squared x/y norms > eps (default 1e-8).
Express CA_j−CA_i in i’s frame and divide by sqrt(local_squared_norm + eps). Concatenate: distance one-hot, pair-valid indicator, right residue one-hot, left residue one-hot, local unit vector(3), pair-valid indicator again.
Zero invalid feature rows, then project with w_features and b_features. Return pseudo_beta, the final pair_mask, raw concatenated features, and learned embedding; zero invalid embeddings after every bias.
Also apply affine population LayerNorm to query with query_scale/query_bias and norm_eps (default 1e-5), project with w_query, broadcast across the template axis, and add to the embedding before its final validity mask.
Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

Why this matters

AlphaFold3 combines template distance bins, residue identities and local directions with the current query-pair representation. Geometry masks determine which template pairs can contribute.

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 af3_template_pair_features(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    eps: float = 1e-08,
    type_count: int = 22,
    norm_eps: float = 1e-05,
    glycine_index: int = 7,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    pseudo_beta: (B, T, N, 3).
    pair_mask: (B, T, N, N).
    features: (B,T,N,N,K+2V+6).
    embedding: (B, T, N, N, Ct)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement af3_template_pair_features")

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.