Species-conditioned self connections and parity-aware gates

Your task

Function signature

equivariant_gated_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    even_gate: str = 'silu',
    residual: bool = False,
) -> dict[str, torch.Tensor]

Worked example

Even and odd gates

import torch

activations = {
    'message': torch.tensor([[[[1.0, 0.0, 0.0]], [[2.0, 4.0, 6.0]], [[2.0, 4.0, 6.0]]]], dtype=torch.float64, requires_grad=True),
    'skip': torch.full((1, 3, 1, 3), 0.0, dtype=torch.float64, requires_grad=True),
    'gates': torch.tensor([[[0.0], [0.0], [1.0]]], dtype=torch.float64, requires_grad=True),
}
parameters = {
    'w_message': torch.tensor([[[1.0]], [[1.0]], [[1.0]]], dtype=torch.float64, requires_grad=True),
    'w_self': torch.tensor([[[[0.0]], [[0.0]], [[0.0]]]], dtype=torch.float64, requires_grad=True),
    'bias': torch.tensor([[0.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'w_gate': torch.tensor([[[1.0]], [[1.0]], [[1.0]]], dtype=torch.float64, requires_grad=True),
    'gate_bias': torch.tensor([[0.0], [0.0], [0.0]], dtype=torch.float64, requires_grad=True),
    'activation_scales': torch.tensor([[[1.0], [1.0], [1.0]], [[1.0], [1.0], [1.0]]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'valid': torch.tensor([True], dtype=torch.bool),
}
structure = {
    'species': torch.tensor([0], dtype=torch.int64),
    'irreps': torch.tensor([[0, 1], [1, -1], [1, -1]], dtype=torch.int64),
    'gate_parities': torch.tensor([1, 1, -1], dtype=torch.int64),
}
even_gate = 'sigmoid'
residual = False

result = equivariant_gated_update(activations=activations, parameters=parameters, masks=masks, structure=structure, even_gate=even_gate, residual=residual)

Expected return

{
    'features': torch.tensor([[[[0.7310585786300049, 0.0, 0.0]], [[1.0, 2.0, 3.0]], [[1.5231883119115297, 3.0463766238230594, 4.569564935734589]]]], dtype=torch.float64),
    'activated_gates': torch.tensor([[[0.0], [0.5], [0.7615941559557649]]], dtype=torch.float64),
    'out_parities': torch.tensor([1, -1, 1], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• For each type t, output copy o and component m, compute z_{ntom}=\sum_c message_{ntcm}\,w^{message}_{tco}+\sum_c skip_{ntcm}\,w^{self}_{species(n),tco}. Weight matrices are input-copy then output-copy; the component axis is unchanged. Add bias only for even scalars (l=0, p=+1).
• For scalar l=0, apply \operatorname{SiLU}(z) if p=+1, \operatorname{tanh}(z) if p=-1, then multiply activation_scales [0,t].
• For l>0, compute a scalar gate g=\mathrm{gates}\mathbin{@}w_{\mathrm{gate}}[t]; add gate_bias only if \mathrm{gate\_parities}[t]=+1. Even gates use config even_gate='silu' default or 'sigmoid'; odd gates always tanh.
• Multiply the gate by activation_scales [1,t], then multiply every component of z by that channel's gate. The output scalar parity stays p; nonscalar output parity is p\cdot \mathrm{gate\_parities}[t].
• Config residual=False by default; True adds raw skip after gating/scalar activation, and is legal only when all nonscalar gate parities are +1 so output irreps match. Reject incompatible residual parity. Return activated_gates zero for scalar types and masked rows.
• Ignore component padding, invalid nodes/species, unused scalar-gate payloads and parity-forbidden bias entries, including NaNs. Global learned weights and activation scales are finite. The scales explicitly expose activation second-moment normalization instead of silently assuming library constants.
• No arbitrary l>0 component bias is ever allowed. 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.
• Use sigmoid(z)=1/(1+e^{-z}) and SiLU(z)=z\,sigmoid(z); odd scalar/gate activations use tanh.

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_gated_update(
    activations: dict,
    parameters: dict,
    masks: dict,
    structure: dict,
    even_gate: str = 'silu',
    residual: bool = False,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    features: (N, T, C, D).
    activated_gates: (N, T, C).
    out_parities: (T,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement equivariant_gated_update")

Sources and further reading

• www.nature.com
• mir-group/nequip · convnetlayer.py
• mir-group/nequip · interaction_block.py
• docs.e3nn.org