Structure module, templates and recycling

Clean Template Recycling

EasyRaw PyTorchAbout 10 min1 public example plus private tests

RFdiffusion

Write and run your solution Sign in

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

Your task

Reuse a cached clean backbone as one template for the next pass. Copy coordinates only when a cache exists and the item has passes remaining; return zeros otherwise.

Function signature

clean_template_recycling(
    activations: dict,
    masks: dict,
    structure: dict,
    max_passes: int = 3,
) -> dict[str, torch.Tensor]

Worked example

Worked example

Item 0 copies its three N/CA/C coordinates into template[0, 0]. Item 1 has no cache and item 2 has finished, so their templates are zero. The new template axis gives shape (3, 1, 1, 3, 3); next_iteration is [2, 1, 3] and continue is [True, True, False].

import torch

activations = {
    'previous_clean': torch.tensor([[[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]]], [[[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [1.0, 1.0, 1.0]]], [[[0.0, 0.0, 2.0], [1.0, 0.0, 2.0], [1.0, 1.0, 2.0]]]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'has_previous': torch.tensor([True, False, True], dtype=torch.bool),
}
structure = {
    'iteration': torch.tensor([1, 0, 3], dtype=torch.int64),
}
max_passes = 3

result = clean_template_recycling(activations=activations, masks=masks, structure=structure, max_passes=max_passes)

Expected return

{
    'template': torch.tensor([[[[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]]]], [[[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]], [[[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]]], dtype=torch.float64),
    'next_iteration': torch.tensor([2, 1, 3], dtype=torch.int64),
    'continue': torch.tensor([True, True, False], dtype=torch.bool),
}

Parameters

ParameterShape and meaning
activationsdictSupplied intermediate tensors.
masksdictSupplied masks. Each field below defines what True includes or permits.
structuredictFields supplied together in this dictionary.
max_passesintDefault 3Maximum recycling passes.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

  • Feedback is active where has_previous is True and iteration < max_passes (default 3). Increment iteration wherever iteration < max_passes, even when no previous prediction is present; continue is next_iteration < max_passes.
  • Return previous_clean as a single template axis, shape [B,1,N,3,3], at active batch entries and exact zero elsewhere. Do not normalize coordinates or update feature states.
  • 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

Recycling feeds information from an earlier prediction into fresh base representations; it does not itself rerun the network.

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 clean_template_recycling(
    activations: dict,
    masks: dict,
    structure: dict,
    max_passes: int = 3,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    template: (B, 1, N, 3, 3).
    next_iteration: (B,).
    continue: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement clean_template_recycling")

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.