Atom-layout gather and masked atom-to-token pooling

Your task

Function signature

atom_layout_gather_and_pool(
    masks: dict,
    atom: torch.Tensor | None = None,
    token: torch.Tensor | None = None,
    atom_to_token: torch.Tensor | None = None,
    source: torch.Tensor | None = None,
    mapping: torch.Tensor | None = None,
    variant: str = 'pool',
) -> dict[str, torch.Tensor]

Worked example

Pool only atoms assigned to usable tokens

import torch

atom = torch.tensor([[[2.0, 0.0], [6.0, 4.0], [10.0, 8.0], [17.0, 31.0], [101.0, 103.0]]], dtype=torch.float64, requires_grad=True)
token = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]], dtype=torch.float64, requires_grad=True)
atom_to_token = torch.tensor([[0, 0, 2, -1, 1]], dtype=torch.int64)
source = None
mapping = None
masks = {
    'atom_mask': torch.tensor([[True, True, True, False, True]], dtype=torch.bool),
    'token_mask': torch.tensor([[True, True, False, True]], dtype=torch.bool),
}
variant = 'pool'

result = atom_layout_gather_and_pool(atom=atom, token=token, atom_to_token=atom_to_token, source=source, mapping=mapping, masks=masks, variant=variant)

Expected return

{
    'broadcast': torch.tensor([[[1.0, 2.0], [1.0, 2.0], [0.0, 0.0], [0.0, 0.0], [3.0, 4.0]]], dtype=torch.float64),
    'pooled': torch.tensor([[[4.0, 2.0], [101.0, 103.0], [0.0, 0.0], [0.0, 0.0]]], dtype=torch.float64),
    'counts': torch.tensor([[2.0, 1.0, 0.0, 0.0]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement two supplied-index representation operations. The only option is variant: 'pool' (default) or 'layout'. Return only that variant's stated keys. All floating inputs are finite, share a dtype (float32 or float64) and device; masks are bool with True meaning present/usable, and index tensors are int64 on the same device.
• B and C are positive. The remaining listed axis lengths may be zero. Inputs may be non-contiguous and must not be modified.
• Pool: atom [B,A,C] contains atom features, token [B,N,C] contains token features, atom_to_token [B,A] supplies a token index for every atom slot, atom_mask [B,A] marks usable atoms and token_mask [B,N] marks usable tokens.
• Let k=\mathrm{atom\_to\_token}[b,a] and let v [b,a] be atom_mask [b,a] AND (0 \le k < N) AND token_mask [b,k]; evaluate the last term only after making indexing safe. Any negative index, including -1, and any index \ge N mean missing/padded, even if atom_mask is True.
• Repeated IDs are allowed, and batch indices never mix. \mathrm{broadcast}[b,a,c] = \mathrm{token}[b,k,c] if v [b,a], otherwise 0. \mathrm{counts}[b,n]=\sum_a\mathbf1[v[b,a]\land k=n]. \mathrm{pooled}[b,n,c] = \sum_{a} 1[v[b,a] \land k=n] \cdot \mathrm{atom}[b,a,c] / \operatorname{max}(\mathrm{counts}[b,n],1).
• Return broadcast [B,A,C], pooled [B,N,C], and counts [B,N]. Counts use atom's floating dtype but are discrete, non-differentiable data. Pooling is a mean over valid atoms, not a sum, and not a residual update. A masked token receives no broadcast contributions, a zero pool and zero count.
• An unoccupied token also has zero pool/count. No \epsilon is used: clamp the count denominator to at least one before division. N=0 is legal and makes every atom invalid; A=0 is also legal.
• Layout: source [B,N,L,C] is a per-token source atom layout, mapping [B,N,M] gives the source slot for each target slot, source_mask [B,N,L] marks present source atoms, and target_mask [B,N,M] marks requested target atoms. Let l=\mathrm{mapping}[b,n,m].
• \mathrm{valid}[b,n,m] = \mathrm{target\_mask}[b,n,m] AND (0 \le l < L) AND source_mask [b,n,l], using safe indexing before reading the source. \mathrm{gathered}[b,n,m,c] = \mathrm{source}[b,n,l,c] if valid, otherwise 0. Return gathered [B,N,M,C] and valid [B,N,M] (bool).
• Negative and out-of-range mappings are padding regardless of target_mask. L and M need not match; maps may repeat or omit slots. This is a gather with zero padding, not pooling or an inverse-map inference. L=0 gives all-zero gathered features and all-false validity; M=0 gives empty outputs.
• The map is supplied, so residue-type lookup tables and atom-name chemistry are outside this contract.

Why this matters

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 atom_layout_gather_and_pool(
    masks: dict,
    atom: torch.Tensor | None = None,  # (B, A, C)
    token: torch.Tensor | None = None,  # (B, N, C)
    atom_to_token: torch.Tensor | None = None,  # (B, A)
    source: torch.Tensor | None = None,  # (B, N, L, C)
    mapping: torch.Tensor | None = None,  # (B, N, M)
    variant: str = 'pool',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    broadcast: (B, A, C). Returned only for variant="pool".
    pooled: (B, N, C). Returned only for variant="pool".
    counts: (B, N). Returned only for variant="pool".
    gathered: (B, N, M, C). Returned only for variant="layout".
    valid: (B, N, M). Returned only for variant="layout"."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement atom_layout_gather_and_pool")

Related exercises

• Builds onAtom–Token Broadcast and Pool
• Builds onAtom Layout Gather

Sources and further reading

• Biohub/esm · layers.py
• google-deepmind/alphafold3 · atom_cross_attention.py
• aqlaboratory/openfold · feats.py