Structure confidence heads

Bidirectional Chain-Pair TM Scores

HardRaw PyTorchAbout 30 min1 public example plus private tests

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

Your task

For each supplied chain pair, consider alignments in either direction and select the best row using the declared weighting rule. Return a symmetric matrix of selected TM scores and support flags in chain_labels order; each entry is selected from the combined directions, not their average.

Function signature

bidirectional_chain_pair_tm_scores(
    pae_probability: torch.Tensor,
    activations: dict,
    masks: dict,
    structure: dict,
    tm_length: str = 'mask',
    count_eps: float = 1e-08,
    tm_interface_length: str = 'global',
    tm_weighted_selection: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Two chains can have different directional evidence

Use two valid tokens in chains [10,20], chain_labels=[10,20], unit token weights, all pairs allowed, tm_length="mask", tm_interface_length="global", and count_eps=0.01. Centers $[0,d_0]$ give TM kernels [1,1/2], with

d_0=1.24\sqrt[3]{19-15}-1.8.

Give diagonal pairs probabilities [1/2,1/2], pair (0,1) probabilities [1/4,3/4], and (1,0) probabilities [3/4,1/4]. The directed chain matrix is

\frac{1}{1.01}\begin{bmatrix}0.75&0.625\\0.875&0.75\end{bmatrix}\approx\begin{bmatrix}0.742574&0.618812\\0.866337&0.742574\end{bmatrix}.

All entries have support. In bidirectional mode, each off-diagonal entry considers both row choices and selects $0.875/1.01\approx0.866337$ ; diagonal entries remain unchanged.

Parameters

Parameter	Shape and meaning
`pae_probability`torch.Tensor	$[B,N,N,L]$ Supplied PAE probabilities. B is batch size, N token count, and L error bins.
`activations`dict	Supplied intermediate tensors.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`tm_length`strDefault `'mask'`	"mask" uses the valid-token count; "weight_floor" uses the sum of valid token weights rounded down for the TM length scale.
`count_eps`floatDefault `1e-08`	Positive offset in the declared TM/confidence weighted-count denominators, added after summing the relevant support.
`tm_interface_length`strDefault `'global'`	"global" uses the overall TM length scale for interfaces; "chain_pair" derives a separate scale from the two participating chains.
`tm_weighted_selection`boolDefault `True`	Use residue weights when selecting the best aligned TM row.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

pae_probability is $[B,N,N,L]$ , with nonnegative bins summing to 1 on valid token pairs; invalid pair values are ignored. pae_centers supplies L nonnegative aligned-error bin centers. weights is a nonnegative importance weight per token; zero invalid token weights.
For a length $n$ , use $d_0=1.24(\max(n,19)-15)^{1/3}-1.8$ and $K_l=d_0^2/(d_0^2+center_l^2)$ . A pair term is the expected K under that pair’s PAE probabilities.
tm_length="mask"(default) uses valid-token count; "weight_floor" uses floor(sum of valid weights). tm_pair_mask additionally restricts scored ordered pairs.
For row i, average pair terms with weights $w_iw_j$ over allowed j, dividing by their sum plus count_eps (default $10^{-8}$ ). Empty rows score zero. tm_weighted_selection=True (default) selects the row maximizing row_score × w_i; False selects the maximum row_score. Return the UNWEIGHTED score of the selected row; ties choose the first row.
The validity flag is True iff at least one allowed pair has positive weights at both endpoints. No support returns zero scores and False, without inventing a positive denominator or fallback.
tm_interface_length="global"(default) uses the global d0 above. "chain_pair" instead uses the sum of the VALID token counts of the endpoint chains (only one chain count for a same-chain pair), before the same d0 clamp. This local length uses token counts even when tm_length="weight_floor".
Return a matrix in the supplied chain_labels order. Each(left,right) entry restricts row endpoints to the left chain and column endpoints to the right chain; the bidirectional mode also includes the reversed ordering. Same-chain diagonal entries are allowed. Apply the same row-selection rule and report support independently per matrix entry.
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

A chain-pair score asks whether two predicted chains fit together. This AF3-style reduction considers alignment frames from both chains before choosing the best supported row, yielding a symmetric score matrix.

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 bidirectional_chain_pair_tm_scores(
    pae_probability: torch.Tensor,  # (B,N,N,L)
    activations: dict,
    masks: dict,
    structure: dict,
    tm_length: str = 'mask',
    count_eps: float = 1e-08,
    tm_interface_length: str = 'global',
    tm_weighted_selection: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    chain_pair_tm: (B, C, C).
    chain_pair_valid: (B, C, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement bidirectional_chain_pair_tm_scores")

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.