One EGNN scalar-feature and coordinate update

Your task

Function signature

egnn_coordinate_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    coord_aggregation: str,
    coord_scale: float,
    residual: bool,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Follow one directed coordinate update

import torch

activations = {
    'h': torch.tensor([[1.0], [2.0]], dtype=torch.float64, requires_grad=True),
    'xyz': torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=torch.float64, requires_grad=True),
    'edge_attr': torch.tensor([[0.0]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'w_edge1': torch.tensor([[0.0], [0.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'b_edge1': torch.tensor([1.0], dtype=torch.float64, requires_grad=True),
    'w_edge2': torch.tensor([[0.0]], dtype=torch.float64, requires_grad=True),
    'b_edge2': torch.tensor([1.0], dtype=torch.float64, requires_grad=True),
    'w_coord': torch.tensor([1.0], dtype=torch.float64, requires_grad=True),
    'w_node1': torch.tensor([[0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'b_node1': torch.tensor([1.0], dtype=torch.float64, requires_grad=True),
    'w_node2': torch.tensor([[0.0]], dtype=torch.float64, requires_grad=True),
    'b_node2': torch.tensor([0.0], dtype=torch.float64, requires_grad=True),
}
masks = {
    'node_valid': torch.tensor([True, True], dtype=torch.bool),
    'edge_valid': torch.tensor([True], dtype=torch.bool),
}
structure = {
    'edge_index': torch.tensor([[1], [0]], dtype=torch.int64),
}
coord_aggregation = 'mean'
coord_scale = 1.0
residual = True

result = egnn_coordinate_update(activations=activations, parameters=parameters, masks=masks, structure=structure, coord_aggregation=coord_aggregation, coord_scale=coord_scale, residual=residual)

Expected return

(torch.tensor([[1.0], [2.0]], dtype=torch.float64), torch.tensor([[0.0, 0.0, 0.0], [1.7310585786300048, 0.0, 0.0]], dtype=torch.float64))

Parameters

Return

tuple[torch.Tensor, torch.Tensor]

Required behavior

• Inputs h [N,C], xyz [N,3], edge_index [2,E], edge_attr [E,A], node_valid [N], edge_valid [E] have N,C,M,\mathrm{Cout}\ge 1, E,A\ge 0. Each column is (\mathrm{receiver} i, \mathrm{sender} j); indices are in [0,N). True means valid.
• An edge is active only when its edge mask and both endpoint masks are True. Directed edges are used exactly as supplied; self edges and repeated edges are allowed and count separately. Mask padded h/xyz and inactive edge_attr before use; padded values may be NaN/Inf.
• For each active edge let d=\mathrm{xyz\_i}-\mathrm{xyz\_j}, q=\operatorname{dot}(d,d), z=\operatorname{concat}(h_{i},h_{j},[q],\mathrm{edge\_attr\_e}), m=\operatorname{SiLU}(\operatorname{SiLU}(z\mathbin{@}w_{\mathrm{edge1}}+b_{\mathrm{edge1}})\mathbin{@}w_{\mathrm{edge2}}+b_{\mathrm{edge2}}). Inactive m is zero.
• Scalar aggregation is ALWAYS M_i=\sum_{e:\,\mathrm{receiver}(e)=i}m_e. Coordinate coefficient a_{e}=\operatorname{dot}(m_{e},w_{\mathrm{coord}}) has no bias.
• Let D_i=\sum_{e:\,\mathrm{receiver}(e)=i}d_e a_e, then divide D_i by \max(\text{number of active incoming edges},1) only if coord_aggregation='mean'; 'sum' leaves it unchanged. coord_scale is any finite scalar and \mathrm{xyz\_next\_i}=\mathrm{xyz\_i}+\mathrm{coord\_scale}\cdot D_{i}.
• No distance normalization, square root, \varepsilon, or stop-gradient is used; coincident coordinates are safe. Compute: u_{i}=\operatorname{SiLU}(\operatorname{concat}(h_{i},M_{i})\mathbin{@}w_{\mathrm{node1}}+b_{\mathrm{node1}})\mathbin{@}w_{\mathrm{node2}}+b_{\mathrm{node2}}.
• If residual=True, require \mathrm{Cout}=C and return h_{\mathrm{next}}=h+u; if False return u. The arguments coord_aggregation, coord_scale, and residual are required.
• Matrix shapes are w_edge1 [2C+1+A,M], w_edge2 [M,M], w_coord [M], w_node1 [C+M,M], w_node2 [M,\mathrm{Cout}], with matching biases b_edge1 [M], b_edge2 [M], b_node1 [M], b_node2 [\mathrm{Cout}]. Return h_next [N,\mathrm{Cout}], xyz_next [N,3], both zero at invalid nodes.
• Empty neighborhoods give no coordinate shift but still apply the node feature MLP. No learned xyz bias, edge attention, dropout, or graph rebuilding is included; strided inputs and empty edge sets are supported.
• Use SiLU(z)=z/(1+e^{-z}) for each named scalar activation.

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 egnn_coordinate_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    coord_aggregation: str,
    coord_scale: float,
    residual: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    h_next: (N, Cout).
    xyz_next: (N, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement egnn_coordinate_update")

Sources and further reading

• vgsatorras/egnn · egnn_clean.py