Protein sequence and inverse folding

Nearest structural-code quantization

EasyRaw PyTorchAbout 10 min2 public examples plus private tests

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

Your task

For each valid latent vector, select the nearest supplied codebook entry by squared distance. Return the selected indices, quantized vectors, and squared reconstruction error.

Function signature

nearest_codebook_quantization(
    latent: torch.Tensor,
    codebook: torch.Tensor,
    mask: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]

Worked example

Compare distances before choosing a code

Use one batch with latent=[(2,1),(1,0)], codebook=[(0,0),(2,0),(0,3)], and mask=[True,True]; there are no settings. For (2,1), squared distances are [5,1,8], so choose code index 1, vector (2,0), and error 1. For (1,0), distances are [1,1,10]; the tie selects index 0, vector (0,0), and error 1. Thus indices=[[1,0]] and squared_error=[[1,1]]; quantized vectors have shape (1,2,2).

Parameters

Parameter	Shape and meaning
`latent`torch.Tensor	$[B, L, C]$ Latent vector to match at each sequence position.
`codebook`torch.Tensor	$[K, C]$ Supplied candidate code vectors; lower index wins exact distance ties.
`mask`torch.Tensor \| NoneDefault `None`	$[B, L]$ Optional bool $[B,L]$ , `True` means valid.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Implement masked nearest-codebook quantization. latent has shape $[B,L,C]$ , codebook $[K,C]$ , with $B,L,C,K\ge 1$ ; all floating tensors share $\mathrm{float32}/64$ dtype and device. Optional mask $[B,L]$ is bool True for a valid token; absent means all valid.
For a valid latent $z$ and code e_k, $d_{k}=\sum_{c} (z_{c}-e_{\mathrm{kc}})^{2}$ . Select the smallest index attaining min_k d_k; do not normalize vectors or distances. Return indices $[B,L]$ int64, $\mathrm{quantized}[B,L,C]=e_{\mathrm{index}}$ , and $\mathrm{squared\_error}[B,L]=d_{\mathrm{index}}$ .
Invalid tokens return index $-1$ , an all-zero vector and exactly zero error, never indexing code $-1$ as real data. No configuration keys are required. Inputs and intermediate squared distances must be finite and representable. Accept noncontiguous latent/codebook layouts and do not mutate inputs.
Gradients are piecewise derivatives with assignments fixed: quantized differentiates through selected code vectors, not latent; squared_error differentiates through both latent and selected codes. Repeated selections accumulate code gradients. Tie outputs obey first-index ordering but tie derivatives are not graded.
Masked tokens contribute zero gradients. This is hard quantization, not a straight-through estimator, EMA codebook training, or a differentiable soft mixture.

Why this matters

Hard assignment is a discrete lookup: selected code vectors receive gradients, while the squared-error output also differentiates with respect to the latent.

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 nearest_codebook_quantization(
    latent: torch.Tensor,  # (B, L, C)
    codebook: torch.Tensor,  # (K, C)
    mask: torch.Tensor | None = None,  # (B, L)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    indices: (B, L).
    quantized: (B, L, C).
    squared_error: (B, L)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement nearest_codebook_quantization")

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.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.