MSA and pair representations

Chain-Aware Relative Position Buckets

EasyRaw PyTorchAbout 10 min1 public example plus private tests

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

Your task

Encode signed residue offsets within a chain, using one extra bucket for pairs from different chains, then project the one-hot buckets.

Function signature

chain_relative_position_buckets(
    residue_index: torch.Tensor,
    mask: torch.Tensor,
    asym_id: torch.Tensor,
    weight: torch.Tensor,
    residue_radius: int = 32,
) -> dict[str, torch.Tensor]

Worked example

Clipping and a cross-chain sentinel

With radius 1, pair (0, 1) has offset 0−2=−2 and clips to bucket 0, selecting weight row [1, 0]. Pair (0, 2) crosses chains, so it uses sentinel bucket 3 and embedding [-2, 4].

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)
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)
residue_radius = 1

result = chain_relative_position_buckets(residue_index=residue_index, mask=mask, asym_id=asym_id, weight=weight, residue_radius=residue_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]$ 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.
`weight`torch.Tensor	$[2r+2,C]$ Learned projection from the ordered one-hot feature blocks to output channels.
`residue_radius`intDefault `32`	Clipping radius for residue-index 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) is a nonnegative integer. 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.
Return only the residue index as the final singleton index axis, its one-hot features of width 2r+2, and the projected embedding.
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 chain_relative_position_buckets(
    residue_index: torch.Tensor,  # (B, N)
    mask: torch.Tensor,  # (B, N)
    asym_id: torch.Tensor,  # (B, N)
    weight: torch.Tensor,  # (2r+2,C)
    residue_radius: int = 32,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    indices: (B,N,N,1).
    features: (B,N,N,2r+2).
    embedding: (B, N, N, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement chain_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.