All-atom folding

Atom Layout Gather

EasyRaw PyTorchAbout 10 min2 public examples plus private tests

AlphaFold2OpenFoldAlphaFold3OpenFold3ESMFold2

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

Your task

Convert between atom layouts using a supplied per-token slot map. Return gathered features and the boolean mask of valid destination slots.

Function signature

gather_atom_layout(
    source: torch.Tensor,
    mapping: torch.Tensor,
    source_mask: torch.Tensor,
    target_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Repeated indices and an empty slot

Destination slots 0 and 1 both select source atom 1, so both receive [5, 7]. Slot 2 has index -1: it receives [0, 0] and False validity. Slot 3 selects source atom 0 and receives [2, 3]. Repeated indices repeat a value; they do not consume it.

import torch

source = torch.tensor([[[[2.0, 3.0], [5.0, 7.0]]]], dtype=torch.float64, requires_grad=True)
mapping = torch.tensor([[[1, 1, -1, 0]]], dtype=torch.int64)
source_mask = torch.tensor([[[True, True]]], dtype=torch.bool)
target_mask = torch.tensor([[[True, True, True, True]]], dtype=torch.bool)

result = gather_atom_layout(source=source, mapping=mapping, source_mask=source_mask, target_mask=target_mask)

Expected return

(torch.tensor([[[[5.0, 7.0], [5.0, 7.0], [0.0, 0.0], [2.0, 3.0]]]], dtype=torch.float64), torch.tensor([[[True, True, False, True]]], dtype=torch.bool))

Parameters

Parameter	Shape and meaning
`source`torch.Tensor	$[B, N, L, C]$ Features in the source atom-slot layout.
`mapping`torch.Tensor	$[B, N, M]$ Target-to-source slot map within each token; out-of-range values mean padding.
`source_mask`torch.Tensor	$[B, N, L]$ True marks a present source atom slot.
`target_mask`torch.Tensor	$[B, N, M]$ True marks a requested target atom slot.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

source is $[B,N,L,C]$ and mapping is $[B,N,M]$ . A target slot is valid only when target_mask is True, its map lies in[0,L), and its mapped source_mask is True.
Gather from the source-slot axis only; tokens and batches never mix. Return zero features in every invalid slot and a boolean valid tensor $[B,N,M]$ .
Repeated source slots may be gathered more than once; gradients accumulate into the reused source feature. Empty source slots are supported through a safe sentinel.

Why this matters

Gather copies features along a known map; pooling reverses the information flow by accumulating atoms into their parent token.

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 gather_atom_layout(
    source: torch.Tensor,  # (B, N, L, C)
    mapping: torch.Tensor,  # (B, N, M)
    source_mask: torch.Tensor,  # (B, N, L)
    target_mask: torch.Tensor,  # (B, N, M)
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    gathered: (B, N, M, C).
    valid: (B, N, M)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement gather_atom_layout")

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.