MSA and pair representations

ESMFold 2 Relative Position Buckets

MediumRaw PyTorchAbout 20 min1 public example plus private tests

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

Your task

Build ESMFold2 relative-position features: residue and token offset buckets, a same-entity indicator, and cross-chain symmetry-copy buckets.

Function signature

esmfold2_relative_position_buckets(
    residue_index: torch.Tensor,
    mask: torch.Tensor,
    asym_id: torch.Tensor,
    token_index: torch.Tensor,
    entity_id: torch.Tensor,
    sym_id: torch.Tensor,
    weight: torch.Tensor,
    residue_radius: int = 32,
    chain_radius: int = 2,
) -> dict[str, torch.Tensor]

Worked example

The two valid tokens have different chains. Pair (0, 1) uses residue/token sentinels and symmetry offset -1, giving indices [3, 3, 0]; the reverse uses [3, 3, 2]. Same-token pairs use [1, 1, 3]. With weight[k]=k, the valid embedding block is [[26, 27], [29, 26]]. Every pair touching token 2 has indices [-1, -1, -1] and zero features/embedding.

import torch

residue_index = torch.tensor([[2, 2, 9]], dtype=torch.int64)
mask = torch.tensor([[True, True, False]], dtype=torch.bool)
asym_id = torch.tensor([[0, 1, 0]], dtype=torch.int64)
token_index = torch.tensor([[0, 0, 0]], dtype=torch.int64)
entity_id = torch.tensor([[4, 4, 4]], dtype=torch.int64)
sym_id = torch.tensor([[0, 1, 0]], dtype=torch.int64)
weight = torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0], [8.0], [9.0], [10.0], [11.0], [12.0]], dtype=torch.float64, requires_grad=True)
residue_radius = 1
chain_radius = 1

result = esmfold2_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, residue_radius=residue_radius, chain_radius=chain_radius)

Expected return

{
    'indices': torch.tensor([[[[1, 1, 3], [3, 3, 0], [-1, -1, -1]], [[3, 3, 2], [1, 1, 3], [-1, -1, -1]], [[-1, -1, -1], [-1, -1, -1], [-1, -1, -1]]]], dtype=torch.int64),
    'features': torch.tensor([[[[0.0, 1.0, 0.0, 0.0, 0.0, 1.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, 0.0, 0.0, 1.0, 1.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, 0.0, 0.0, 0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 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, 0.0, 0.0]], [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.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([[[[26.0], [27.0], [0.0]], [[29.0], [26.0], [0.0]], [[0.0], [0.0], [0.0]]]], dtype=torch.float64),
}

Parameters

Parameter	Shape and meaning
`residue_index`torch.Tensor	$[B, N]$ Supplied residue numbering, not required to start at zero.
`mask`torch.Tensor	$[B, N]$ True for valid tokens.
`asym_id`torch.Tensor	$[B, N]$ Chain identity per token; equal IDs mean the same chain.
`token_index`torch.Tensor	$[B, N]$ Token number within the supplied residue convention; multiple tokens may share a residue.
`entity_id`torch.Tensor	$[B, N]$ Chemical equivalence identity; equal IDs set the same-entity indicator.
`sym_id`torch.Tensor	$[B, N]$ Symmetry-copy index used for cross-chain offsets.
`weight`torch.Tensor	$[4r+2c+7,C]$ Learned projection from the ordered one-hot feature blocks to output channels.
`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

residue_index and mask have shape $[B,N]$ , with True marking a valid token. residue_radius r (default 32) and chain_radius c (default 2, when used) are nonnegative integers. Weights project the concatenated bucket features into embedding channels.
Same-chain pairs have equal asym_id. Their residue bucket is clamp(residue_index[i]−residue_index[j]+r,0,2r); other chains use sentinel 2r+1.
Invalid endpoint pairs have indices−1, all-zero features, and all-zero embeddings.
The token bucket uses token_index[i]−token_index[j] clipped/shifted by r only for same-chain pairs at the same residue_index. Every other pair uses sentinel 2r+1.
Append same_entity as one scalar indicator. The symmetry bucket uses sym_id[i]−sym_id[j] clipped/shifted by c when asym_id differs; otherwise use sentinel 2c+1.
Feature order is residue one-hot(2r+2), token one-hot(2r+2), same_entity(1), symmetry one-hot(2c+2). Return the three corresponding integer indices in residue/token/symmetry order.
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

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 esmfold2_relative_position_buckets(
    residue_index: torch.Tensor,  # (B, N)
    mask: torch.Tensor,  # (B, N)
    asym_id: torch.Tensor,  # (B, N)
    token_index: torch.Tensor,  # (B, N)
    entity_id: torch.Tensor,  # (B, N)
    sym_id: torch.Tensor,  # (B, N)
    weight: torch.Tensor,  # (4r+2c+7,C)
    residue_radius: int = 32,
    chain_radius: int = 2,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    indices: (B,N,N,3).
    features: (B,N,N,4r+2c+7).
    embedding: (B, N, N, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement esmfold2_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.