Structure confidence heads

Resolved-Atom Head

EasyRaw PyTorchAbout 10 min3 public examples plus private tests

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

Your task

Predict two logits per atom slot: unresolved and resolved. Return logits only; probability conversion is a separate step.

Function signature

resolved_atom_head(
    activations: dict,
    parameters: dict,
    masks: dict,
) -> torch.Tensor

Worked example

An unresolved/resolved affine readout

For one present token and atom, use single=[1,3], W_resolved[:,0,:]=[[1,2],[3,-1]], and b_resolved[0,:]=[0.5,-0.5], with both masks true. The unresolved logit is 1·1+3·3+0.5=10.5; the resolved logit is 1·2+3·(-1)-0.5=-1.5. Return these two scores in that order with shape (1,1,1,2). They are logits, so negative values are valid.

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.

Return

torch.Tensor · $[B, N, A, 2]$

Return Unresolved/resolved atom logits.

Required behavior

single is $[B,N,C]$ ; token_mask and atom_mask mark included token/atom slots. Sanitize absent token features before the learned head.
Use the single features directly, without normalization or an activation.
Contract W_resolved $[C,A,2]$ with the single feature axis and add b_resolved $[A,2]$ . The final class order is unresolved, resolved. Return shape $[B,N,A,2]$ , with absent token/atom logits exactly zero after bias.
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 resolved atom is experimentally observed. This head predicts unresolved/resolved logits separately from chemical existence; input masks say which atom slots can receive predictions.

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 resolved_atom_head(
    activations: dict,
    parameters: dict,
    masks: dict,
) -> torch.Tensor:
    """Return a tensor with shape (B, N, A, 2)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement resolved_atom_head")

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.