MSA and pair representations · INTEGRATION

Chain, entity and token-relative position buckets

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

ProteinMPNN AlphaFold2 AlphaFold3 ESMFold ESMFold2 ESM3

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

Your task

Represent each ordered token pair by categories describing its signed residue/token separation and the applicable chain, entity, or symmetry relationship. Each embedding[b,i,j] is the supplied projection of those category features. Implement the selected convention, including its exact direction and padding behavior.

Function signature

relative_position_buckets(
    residue_index: torch.Tensor,
    mask: torch.Tensor,
    weight: torch.Tensor,
    asym_id: torch.Tensor | None = None,
    token_index: torch.Tensor | None = None,
    entity_id: torch.Tensor | None = None,
    sym_id: torch.Tensor | None = None,
    variant: str = 'chain',
    residue_radius: int = 32,
    chain_radius: int = 2,
) -> dict[str, torch.Tensor]

Worked example

Clipped offsets and cross-chain categories

In the chain example, radius 1 makes pair (0,1), with residue offset 0−2=−2, clip to bucket 0 and select weight row [1,0]. Pair (0,2) crosses chains, so it uses sentinel bucket 3 and embedding [-2,4]. The selected variant’s rules determine whether a masked pair is zero or uses ESMFold’s learned padding category.

import torch

residue_index = torch.tensor([[0, 2, 0]], dtype=torch.int64)
mask = torch.tensor([[True, True, True]], dtype=torch.bool)
asym_id = torch.tensor([[0, 0, 1]], dtype=torch.int64)
token_index = None
entity_id = None
sym_id = None
weight = torch.tensor([[1.0, 0.0], [0.0, 2.0], [3.0, -1.0], [-2.0, 4.0]], dtype=torch.float64, requires_grad=True)
variant = 'chain'
residue_radius = 1
chain_radius = 2

result = relative_position_buckets(residue_index=residue_index, mask=mask, asym_id=asym_id, token_index=token_index, entity_id=entity_id, sym_id=sym_id, weight=weight, variant=variant, residue_radius=residue_radius, chain_radius=chain_radius)

Expected return

{
    'indices': torch.tensor([[[[1], [0], [3]], [[2], [1], [3]], [[3], [3], [1]]]], dtype=torch.int64),
    'features': torch.tensor([[[[0.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]], [[0.0, 0.0, 1.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]], [[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0]]]], dtype=torch.float64),
    'embedding': torch.tensor([[[[0.0, 2.0], [1.0, 0.0], [-2.0, 4.0]], [[3.0, -1.0], [0.0, 2.0], [-2.0, 4.0]], [[-2.0, 4.0], [-2.0, 4.0], [0.0, 2.0]]]], dtype=torch.float64),
}

Parameters

Parameter	Shape and meaning
`residue_index`torch.Tensor	$[B, N]$ `int64` $[B,N]$ ; supplied residue numbering, not required to start at zero.
`mask`torch.Tensor	$[B, N]$ `bool` $[B,N]$ , `True` for valid tokens.
`asym_id`torch.Tensor \| NoneDefault `None`	$[B, N]$ `int64` $[B,N]$ chain identities; required except for esmfold.
`token_index`torch.Tensor \| NoneDefault `None`	$[B, N]$ `int64` $[B,N]$ , required for $\mathrm{af3}/\mathrm{esmfold2}$ ; multiple tokens may share a residue.
`entity_id`torch.Tensor \| NoneDefault `None`	$[B, N]$ `int64` $[B,N]$ chemical/entity equivalence classes; $\mathrm{af3}/\mathrm{esmfold2}$ only.
`sym_id`torch.Tensor \| NoneDefault `None`	$[B, N]$ `int64` $[B,N]$ symmetry-copy `indices`; $\mathrm{af3}/\mathrm{esmfold2}$ only.
`weight`torch.Tensor	$[F, C]$ $\mathrm{float32}/64[F,C]$ ; $F$ is determined by variant and radii; $C\ge 1$ .
`variant`strDefault `'chain'`	Select chain-aware residue offsets, AlphaFold3 token/entity features, ESMFold2 token/entity features, or ESMFold padding-aware offsets.
`residue_radius`intDefault `32`	Clipping radius for residue-index differences.
`chain_radius`intDefault `2`	Clipping radius for symmetry-copy differences.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Build discrete relative-position features and apply a supplied feature projection. Inputs use batch $B$ and token $N$ axes, with $B$ , $N \ge 1$ . All identifiers are int64, differences fit int64, mask $[B,N]$ is Boolean True for a valid token, and weight $[F,C]$ is $\mathrm{float32}/64$ .
Return indices $[B,N,N,G]$ (int64), features $[B,N,N,F]$ and $\mathrm{embedding}=\mathrm{features}\mathbin{@}\mathrm{weight}[B,N,N,C]$ . Matrix row $i$ always names the first token.
variant defaults to 'chain'; residue_radius $R$ defaults to $32$ and chain_radius $S$ to $2$ ; both are nonnegative integers. Define $\mathrm{clip\_R}(d)=\operatorname{min}(2R,\operatorname{max}(0,d+R))$ .
* chain (ProteinMPNN): residue $\mathrm{bucket}=\mathrm{clip\_R}(\mathrm{residue\_i}-\mathrm{residue\_j})$ if $\mathrm{asym\_i}=\mathrm{asym\_j}$ , else $2R+1$ . One-hot width $F=2R+2$ and $G=1$ .
* af3 and esmfold2: concatenate residue one-hot (same rule), token one-hot ( $\mathrm{clip\_R}(\mathrm{token\_i}-\mathrm{token\_j})$ only when both chain AND residue IDs match, else $2R+1$ ), one same-entity scalar, and symmetry one-hot of width $2S+2$ .
AF3 keeps $\mathrm{clip\_S}(\mathrm{sym\_i}-\mathrm{sym\_j})$ when $\mathrm{entity\_i}=\mathrm{entity\_j}$ , otherwise uses $2S+1$ . ESMFold2 instead keeps it when $\mathrm{asym\_i}\ne \mathrm{asym\_j}$ , otherwise uses $2S+1$ . These source conventions are deliberately different.
$F=2\cdot (2R+2)+1+(2S+2)$ ; $G=3$ stores residue, token and symmetry buckets, not the entity scalar. * esmfold (v1): $\mathrm{bucket}=\operatorname{clamp}(\mathrm{residue\_j}-\mathrm{residue\_i},-R,R)+R+1$ ; an invalid pair uses bucket $0$ . $F=2R+2$ , $G=1$ .
Its reserved padding bucket is a REAL one-hot, so embedding includes weight $[0]$ even for invalid pairs, matching the source lookup. ESM3's structural tokenizer uses this same formula for its query residue against a supplied neighborhood: select that query's row with all neighborhood entries valid.
Constructing the neighborhood itself is separate. AlphaFold2 monomer's relative one-hot is reproduced by chain with all asym_id equal and the last, unused feature column removed: its native width is $2R+1$ , not $2R+2$ . A native bias-free projection can use the first $2R+1$ weight rows plus an unused final row.
This does not claim an AF2 interchain feature; source affine readouts with bias can be composed using affine_projection.
For $\mathrm{chain}/\mathrm{af3}/\mathrm{esmfold2}$ an invalid pair (either token masked) has all $\mathrm{indices}=-1$ and exact-zero features/embedding.
For all variants integer metadata at masked tokens is ignored. One-hots are constructed from comparisons, not torch.nn.functional.one_hot. No bias, tokenization or sequence alignment is performed.
Invalid variant or negative radius raises ValueError; incompatible shapes, nonfinite weights, overflowing identifier differences and projection products/sums outside the common dtype's finite range are outside the contract.

Why this matters

Relative offsets become discrete feature channels, but source models differ in signs, special buckets and which identity comparisons enable each feature.

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 relative_position_buckets(
    residue_index: torch.Tensor,  # (B, N)
    mask: torch.Tensor,  # (B, N)
    weight: torch.Tensor,  # (F, C)
    asym_id: torch.Tensor | None = None,  # (B, N)
    token_index: torch.Tensor | None = None,  # (B, N)
    entity_id: torch.Tensor | None = None,  # (B, N)
    sym_id: torch.Tensor | None = None,  # (B, N)
    variant: str = 'chain',
    residue_radius: int = 32,
    chain_radius: int = 2,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    indices: (B, N, N, G).
    features: (B, N, N, F).
    embedding: (B, N, N, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement relative_position_buckets")

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.