Atom–Token Broadcast and Pool

Your task

Function signature

atom_token_pool(
    atom: torch.Tensor,
    token: torch.Tensor,
    atom_to_token: torch.Tensor,
    atom_mask: torch.Tensor,
    token_mask: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Repeated atoms and a masked token

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)
atom_mask = torch.tensor([[True, True, True, False, True]], dtype=torch.bool)
token_mask = torch.tensor([[True, True, False, True]], dtype=torch.bool)

result = atom_token_pool(atom=atom, token=token, atom_to_token=atom_to_token, atom_mask=atom_mask, token_mask=token_mask)

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

• atom has shape (B,A,C) and token has shape (B,N,C). An atom contributes only when atom_mask is True, its map lies in [0,N), and the mapped token_mask is True.
• Broadcast the mapped token vector to each contributing atom; all other atom vectors are zero. Pool contributing atom vectors by exact sum/count; tokens with no atoms return zero.
• counts has shape (B,N) in the floating feature dtype. Repeated token IDs accumulate contributions and gradients. Sentinel or out-of-range IDs contribute nothing. A zero-length token axis is supported.
• Use supplied CPU torch.float32 or torch.float64 tensors with shared floating dtype/device. Do not mutate inputs or initialize new parameters; preserve arbitrary valid leading dimensions and accept noncontiguous inputs.

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_token_pool(
    atom: torch.Tensor,  # (B, A, C)
    token: torch.Tensor,  # (B, N, C)
    atom_to_token: torch.Tensor,  # (B, A)
    atom_mask: torch.Tensor,  # (B, A)
    token_mask: torch.Tensor,  # (B, N)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    broadcast: (B, A, C).
    pooled: (B, N, C).
    counts: (B, N)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement atom_token_pool")

Related exercises

• Builds onAtom Layout Gather
• Builds onMasked Mean
• Combined exerciseAtom-layout gather and masked atom-to-token pooling

Sources and further reading

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