Radial-weighted irrep convolution with receiver aggregation

Your task

Function signature

equivariant_message_passing(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    normalization: str = 'sender_sqrt',
    radial_activation: str = 'silu',
    symmetry: str = 'o3',
) -> dict[str, torch.Tensor]

Worked example

A vector message uses sender normalization and a radial weight

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Clean node component padding and node masks, then h[n,t,o,m]=\sum_{c} \mathrm{features}[n,t,c,m]\cdot w_{\mathrm{in}}[t,c,o]. Config normalization='sender_sqrt' default divides h [n] by \sqrt{\mathrm{avg\_neighbors}[n]} before gathering the sender; this is the audited NequIP placement.
• Alternatives 'sum' and 'mean' do not scale h; mean divides each receiver's final sum by its accepted incoming-edge count clamped to at least one. An edge is accepted only if edge_valid and both receiver/sender are valid. False edge_valid indices may be sentinel values.
• Compute radial weights W=\mathrm{activation}(\mathrm{radial}\mathbin{@}\mathrm{radial\_w1}+\mathrm{radial\_b1})\mathbin{@}\mathrm{radial\_w2}+\mathrm{radial\_b2}, reshaped [E,P,\mathrm{Ci},\mathrm{Co}]; radial_activation='silu' default or 'relu'.
• For each legal CG path (input type, edge harmonic type, output type), couple h [\mathrm{sender}] and edge_sh with its fixed CG tensor, then channel-mix by W.
• Multiply each path by \mathrm{path\_scale}/\sqrt{\mathrm{Ci}\cdot \mathrm{active\_paths\_to\_output\_type}}, sum paths by output type, apply type-wise w_out on channels, multiply edge_weight, and sum into receiver. This is algebraically the same as applying w_out after aggregation.
• symmetry='o3' default enforces triangle bounds and \mathrm{parity\_out}=\mathrm{parity\_in}\cdot \mathrm{parity\_edge}; 'so3' omits parity filtering. Return also final per-edge messages, accepted degree, and W zeroed for rejected edges/illegal paths. No tensor-component bias is allowed.
• All false-edge payloads, invalid node data/avg_neighbors and padded components may be nonfinite and must be ignored. E=0 is supported; N\ge 1. A packed block (l,p) has 2\cdot l+1 real components and multiplicity C; p is +1 or -1 under inversion. D is a common padded width.
• Padding is ignored, including NaN padding. Distinct copies of a type are channels, not XYZ coordinates. The fixture basis is Cartesian l=1 [x,y,z] and l=2 orthonormal symmetric-traceless tensors:
• Q=[(\mathrm{xz}+\mathrm{zx})/\sqrt{2},(\mathrm{xy}+\mathrm{yx})/\sqrt{2},\mathrm{diag}(-1,2,-1)/\sqrt{6},(\mathrm{yz}+\mathrm{zy})/\sqrt{2},\mathrm{diag}(-1,0,1)/\sqrt{2}]. Fixed row-orthonormal CG tensors intertwine these bases; phases are defined by the supplied tables, not inferred from an external checkpoint.
• Supplied higher-degree tables must obey the same stated normalization and representation law.

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 equivariant_message_passing(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    normalization: str = 'sender_sqrt',
    radial_activation: str = 'silu',
    symmetry: str = 'o3',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    features: (N, O, Co, D).
    edge_messages: (E, O, Co, D).
    radial_weights: (E, P, Ci, Co).
    degree: (N,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement equivariant_message_passing")

Sources and further reading

• www.nature.com
• mir-group/nequip · interaction_block.py
• mir-group/nequip · _tp_scatter_base.py
• arxiv.org
• gcorso/DiffDock · tensor_layers.py