Structure module, templates and recycling

Template Pair Mean

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 residue pair, average its valid template feature vectors using the stated count-plus-epsilon denominator. Apply ReLU to that mean and then the supplied affine projection. Return the resulting pair update and each template’s averaging coefficient; pairs with no valid template stay zero after the bias.

Function signature

template_pair_mean(
    activations: dict,
    parameters: dict,
    masks: dict,
    mean_eps: float = 1e-07,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Average before applying ReLU

For one valid residue pair, take two valid template vectors [-2,2] and [4,4], set mean_eps=0, and use an identity output matrix with zero bias. The weights are [1/2,1/2], so the mean is [1,3]; ReLU and the projection return [1,3]. Applying ReLU separately before averaging would instead give ([0,2]+[4,4])/2=[2,3], which is a different operation.

Parameters

Parameter	Shape and meaning
`activations`dict	Supplied intermediate tensors.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`mean_eps`floatDefault `1e-07`	Nonnegative value added once to the usable-template count; empty pairs return zero.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

template is $[B,T,N,N,Ct]$ ; combine template_mask with pair_mask broadcast across T. True means included. Invalid template payloads must not influence projections or sums.
Compute the masked sum over T divided by count + mean_eps (default $10^{-7}$ , nonnegative). Empty pairs have exact zero mean. Return weights $[B,T,N,N]$ equal to mask/(count + mean_eps) on nonempty pairs.
Apply ReLU to the masked template mean, then w_mean and b_mean. There is no query, attention projection, gate, or residual.
After the final bias, zero every pair with no valid template. Template count T is positive; all-masked pairs are supported.

Why this matters

Templates offer alternative descriptions of the same residue pair. Attention lets the query choose among them; mean mode treats usable templates equally before a ReLU and learned readout.

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 template_pair_mean(
    activations: dict,
    parameters: dict,
    masks: dict,
    mean_eps: float = 1e-07,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    update: (B,N,N,C_q).
    weights: (B,T,N,N)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement template_pair_mean")

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.