Nonbacktracking directed-bond recurrence and atom readout

Your task

Function signature

nonbacktracking_bond_messages(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    depth: int = 3,
    undirected: bool = False,
    dropout_p: float = 0.0,
) -> dict[str, torch.Tensor]

Worked example

Trace a three-atom chain

import torch

activations = {
    'atom_features': torch.tensor([[1.0], [1.0], [1.0]], dtype=torch.float64, requires_grad=True),
    'edge_features': torch.tensor([[0.0], [0.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'w_i': torch.tensor([[1.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'b_i': torch.tensor([0.0], dtype=torch.float64, requires_grad=True),
    'w_h': torch.tensor([[1.0]], dtype=torch.float64, requires_grad=True),
    'b_h': torch.tensor([0.0], dtype=torch.float64, requires_grad=True),
    'w_o': torch.tensor([[1.0], [1.0]], dtype=torch.float64, requires_grad=True),
    'b_o': torch.tensor([0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'atom_mask': torch.tensor([True, True, True], dtype=torch.bool),
    'edge_mask': torch.tensor([True, True, True, True], dtype=torch.bool),
    'message_keep': torch.tensor([[[True], [True], [True], [True]]], dtype=torch.bool),
    'atom_keep': torch.tensor([[True], [True], [True]], dtype=torch.bool),
}
structure = {
    'edge_index': torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.int64),
    'reverse_index': torch.tensor([1, 0, 3, 2], dtype=torch.int64),
}
depth = 2
undirected = False
dropout_p = 0.0

result = nonbacktracking_bond_messages(activations=activations, parameters=parameters, masks=masks, structure=structure, depth=depth, undirected=undirected, dropout_p=dropout_p)

Expected return

{
    'initial_projection': torch.tensor([[1.0], [1.0], [1.0], [1.0]], dtype=torch.float64),
    'bond_states': torch.tensor([[[1.0], [1.0], [1.0], [1.0]], [[1.0], [2.0], [2.0], [1.0]]], dtype=torch.float64),
    'atom_messages': torch.tensor([[2.0], [2.0], [2.0]], dtype=torch.float64),
    'atom_hidden': torch.tensor([[3.0], [3.0], [3.0]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Implement the directed-bond recurrence and bond-to-atom readout; do not replace it with a generic node MLP. Let X=\mathrm{atom\_features}[N,A], E=\mathrm{edge\_features}[D,B], with A,B,H\ge 1 and N,D possibly zero. Floating tensors share a finite CPU \mathrm{float32}/\mathrm{float64} dtype.
• edge_index [2,D] int64 stores source s(e) then destination t(e). reverse_index [D] int64 maps each active edge to its unique paired reverse r(e): r(r(e))=e, r(e)\ne e, s(r(e))=t(e), t(r(e))=s(e). Active edges have distinct in-range endpoints whose atom_mask is True, and their reverse is active.
• Parallel edge pairs are allowed; exclude the paired identity, not every edge with the same endpoints. edge_mask [D] and atom_mask [N] are bool, True means valid. Inactive edge and reverse-index payloads may be arbitrary integers and must never be indexed.
• Invalid graph invariants other than the documented configuration checks are outside the contract. Features at masked rows are ignored. All projections use row vectors: w_i [A+B,H], b_i [H], w_h [H,H], b_h [H], w_o [A+H,H], b_o [H].
• Define \rho(z)=\operatorname{max}(z,0) elementwise with derivative zero at z=0.
• For valid edges: \begin{aligned}Z_e&=\operatorname{concat}(X_{s(e)},E_e)\mathbin{@}w_i+b_i,\\H_e^0&=\rho(Z_e).\end{aligned} Initialization is not dropped out. Preserve Z, the preactivation projection, as the residual at EVERY update. depth defaults to 3 and must be an integer from 1 through 8.
• Execute exactly \mathrm{depth}-1 updates. undirected defaults to False. At each update t=1,...,\mathrm{depth}-1 let G=H^{t-1} in directed mode; in the supported undirected option, first use G_{e}=(H_{e}^{t-1}+H_{r(e)}^{t-1})/2. This does not average Z and does not overwrite the previously returned state.
• At each update, compute M_e=\sum_{f\ active:\,t(f)=s(e),\ f\ne r(e)}G_f, then H_e^t=ReLU(Z_e+M_e w_h+b_h)\odot message\_keep[t-1,e]/(1-p). The supplied keep mask has shape (depth−1,D,H); p=dropout_p lies in [0,1) and defaults to zero.
• The deterministic mask-and-scale operation is always applied; for ordinary no-dropout evaluation use p=0 and all-True masks. Masks are treated as given discrete outcomes, never sampled. After the last state H^{\mathrm{depth}-1}, sum all incoming active directed bonds (no reverse subtraction here):
• \begin{aligned}A_v&=\sum_{\substack{e\ \mathrm{active}\\t(e)=v}}H_e^{\mathrm{depth}-1},\\Y_v&=\rho\bigl(\operatorname{concat}(X_v,A_v)\mathbin{@}w_o+b_o\bigr)\odot\frac{\mathrm{atom\_keep}[v]}{1-p}.\end{aligned} Here atom_keep [N,H] is also supplied Boolean.
• Return \mathrm{initial\_projection}=Z[D,H], bond_states [\mathrm{depth},D,H] containing H^{0} through H^{\mathrm{depth}-1}, \mathrm{atom\_messages}=A[N,H], and \mathrm{atom\_hidden}=Y[N,H]. Mask Z and every bond state to exactly zero for inactive edges, and atom outputs to zero for padded atoms.
• An isolated valid atom has A_{v}=0 but may have a nonzero descriptor/bias readout. Empty axes are retained, not replaced by dummy nodes. All outputs remain connected with zero derivatives to unused parameters, including w_h/b_h at \mathrm{depth}=1 and edge-only parameters when D=0. Reject invalid depth/dropout_p with ValueError.
• Optional additional atom-descriptor projections, molecular pooling, task prediction heads and training are outside this core; no descriptors, weights, random variates or chemical structures are inferred or loaded.

Why this matters

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 nonbacktracking_bond_messages(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    depth: int = 3,
    undirected: bool = False,
    dropout_p: float = 0.0,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    initial_projection: (D, H).
    bond_states: (depth, D, H).
    atom_messages: (N, H).
    atom_hidden: (N, H)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement nonbacktracking_bond_messages")

Sources and further reading

• doi.org
• chemprop/chemprop · mixins.py
• chemprop/chemprop · base.py