Irrep-safe batch and instance normalization

Your task

Function signature

equivariant_batch_normalization(
    x: torch.Tensor,
    valid: torch.Tensor,
    irreps: torch.Tensor,
    multiplicities: torch.Tensor,
    running_mean: torch.Tensor,
    running_var: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    options: dict,
) -> dict[str, torch.Tensor]

Worked example

Center scalars, scale whole vectors

import torch

x = torch.tensor([[[[[1.0, 0.0, 0.0]], [[1.0, 0.0, 0.0]]]], [[[[3.0, 0.0, 0.0]], [[0.0, 1.0, 0.0]]]]], dtype=torch.float64, requires_grad=True)
valid = torch.tensor([[True], [True]], dtype=torch.bool)
irreps = torch.tensor([[0, 1], [1, -1]], dtype=torch.int64)
multiplicities = torch.tensor([1, 1], dtype=torch.int64)
running_mean = torch.tensor([[0.0], [0.0]], dtype=torch.float64, requires_grad=True)
running_var = torch.tensor([[1.0], [1.0]], dtype=torch.float64, requires_grad=True)
weight = torch.tensor([[1.0], [1.0]], dtype=torch.float64, requires_grad=True)
bias = torch.tensor([[0.0], [0.0]], dtype=torch.float64, requires_grad=True)
options = {
    'eps': 1.0,
    'momentum': 0.5,
}

result = equivariant_batch_normalization(x=x, valid=valid, irreps=irreps, multiplicities=multiplicities, running_mean=running_mean, running_var=running_var, weight=weight, bias=bias, options=options)

Expected return

{
    'output': torch.tensor([[[[[-0.7071067811865475, 0.0, 0.0]], [[0.8660254037844387, 0.0, 0.0]]]], [[[[0.7071067811865475, 0.0, 0.0]], [[0.0, 0.8660254037844387, 0.0]]]]], dtype=torch.float64),
    'mean': torch.tensor([[[2.0], [0.0]]], dtype=torch.float64),
    'second_moment': torch.tensor([[[1.0], [0.3333333333333333]]], dtype=torch.float64),
    'running_mean': torch.tensor([[1.0], [0.0]], dtype=torch.float64),
    'running_var': torch.tensor([[1.0], [0.6666666666666666]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs x [B,S,T,C,D] pack B outer instances, S sample positions, T real irrep types, at most C multiplicity copies per type, and padded component width D. Each \mathrm{irreps}[t]=(l,p), l\ge 0,p in {-1,+1}; d_{t}=2\cdot l+1\le D. multiplicities [t] is an integer in [0,C].
• Copy c is active iff c<\mathrm{multiplicities}[t]; component m is active iff m<d_{t}. valid [B,S] uses True for observed positions. B or S may be zero; T,C,D\ge 1. Active observed values, active running state and enabled active affine parameters are finite.
• Required centered squares, reductions, \epsilon-shifted moments, gains and final affine arithmetic must also be representable and finite in the input dtype: arbitrary finite float32 inputs near 10^{20} whose squares overflow are outside the numerical domain.
• Nonfinite values at masked positions, padded copies/components, forbidden biases or nonexistent scalar-mean slots are ignored before arithmetic. All weight/bias entries are also ignored before arithmetic when affine=False, even if nonfinite. Active running state still must be finite in instance mode because it is returned unchanged.
• The irrep basis must be orthonormal. No per-component affine parameters are allowed. defaults match original DiffDock's pinned e3nn 0.5.0 BatchNorm: mode='train', instance=False, normalization='component', reduce='mean', affine=True, \varepsilon=10^{-5}, \mathrm{momentum}=0.1.
• mode may be 'train'/'eval'; normalization 'component'/'norm'; reduce 'mean'/'max'; \varepsilon>0 and 0\le \mathrm{momentum}\le 1 are finite. Floating tensors use float32 or float64.
• Eps must remain finite and at least the smallest positive NORMAL number in x.dtype (\mathrm{float32}:1.1754943508222875\times 10^{-38}, \mathrm{float64}:2.2250738585072014\times 10^{-308}).
• This rejects underflowing and subnormal \varepsilon, such as float32 \varepsilon 10^{-50} or 10^{-40}, whose normalization backward can overflow even for zero features. As with forward arithmetic, VJP claims require finite representable first-order intermediates and gradients for the supplied cotangents.
• Source DiffDock passes [N_{\mathrm{nodes}},\mathrm{irreps.dim}], equivalent here to B=N_{\mathrm{nodes}},S=1, so default statistics mix nodes across complexes. B does not automatically mean molecular graph. The optional instance=True variant uses each B row independently and is not DiffDock's default.
• ONLY l=0,p=+1 ('0\mathrm{e}') is a scalar for centering and bias. The odd l=0 pseudoscalar ('0\mathrm{o}') is NEVER centered or biased.
• For current statistics (mode='train' OR instance=True), compute each active 0\mathrm{e} copy's mean \mu: global mode pools all valid (b,s) observations equally; instance mode averages valid s within that b. All other mean slots are zero. Empty statistics have mean zero.
• Global eval instead uses supplied running_mean for active 0\mathrm{e} copies. Subtract \mu from 0\mathrm{e} fields only, leaving odd scalars/vectors/higher-degree tensors uncentered; sanitize masks and component padding again.
• For current statistics, compute energy e[b,s,t,c]=\sum_{m<d_{t}} \mathrm{centered\_x}^{2}.
• In normalization='component', divide by d_t, not padded D; 'norm' uses the component sum. First reduce over valid S positions independently per B row: 'mean' averages valid energies, 'max' chooses the largest valid energy (first valid position at a tie). Empty instances have zero and do not participate in global averaging.
• instance=True uses each instance's statistic directly. Otherwise average these per-instance statistics EQUALLY over nonempty B rows, even if their valid counts differ. This is the source's sample-then-batch reduction hierarchy with an explicit ragged-mask extension.
• Scalar centering still pools individual valid observations globally, so under ragged masks its weights can differ from the final norm-statistic weights.
• For example counts 3 and 1 with 0\mathrm{e} values (0,0,0) and (4) give global \mu=1 and \mathrm{second\_moment}=(1+9)/2=5, not pooled variance 3. Global eval uses supplied running_var instead.
• This stored quantity is a biased variance for centered 0\mathrm{e} and a second moment/RMS energy for all other irreps; there is no n/(n-1) correction.
• For each full active irrep copy, output \mathrm{centered\_x}/\sqrt{\mathrm{second\_moment}+\varepsilon}, multiplied by weight [t,c] if affine=True. Only active 0\mathrm{e} copies then add bias [t,c]. If affine=False ignore supplied weight/bias, with defined zero derivatives for them.
• One copy's gain is shared across all d_t components. Invalid positions and padded copies/components return zero, including after a nonzero scalar bias. There is no nonlinearity or cross-copy mixing.
• Global train with any observed position returns functional next running statistics: \mathrm{running\_mean}'=(1-\mathrm{momentum})\,\mathrm{old\_mean}+\mathrm{momentum}\,\mu for active 0\mathrm{e} copies, and \mathrm{running\_var}'=(1-\mathrm{momentum})\,\mathrm{old\_var}+\mathrm{momentum}\,\mathrm{second\_moment} for every active copy.
• With no observations, preserve supplied active running state. Eval OR instance mode leaves supplied active state unchanged; e3nn instance mode has no state, and this API's pass-through tensors merely keep a uniform functional interface. Nonexistent mean slots and padded state slots return zero.
• Active running_var is finite and nonnegative, even when it is unused for normalizing the current batch. Source e3nn detaches buffer updates; this exercise instead returns DIFFERENTIABLE pure-functional EMA tensors, explicitly grading their VJPs as a teaching adaptation. Never mutate inputs or detach values.
• Return output [B,S,T,C,D], mean [G,T,C], second_moment [G,T,C], running_mean [T,C], running_var [T,C], where G=1 for global and G=B for instance mode. Mean is zero outside 0\mathrm{e}, and both diagnostics are zero at padded copies.
• Empty current statistics are zero; global eval diagnostics still describe the supplied running state. Max-mode VJPs are promised away from meaningful winner ties; zero-energy ties have zero first derivative through the square.
• Unsupported configuration or negative active running variance raises ValueError; invalid irrep/multiplicity shapes or nonorthonormal bases are outside contract.

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_batch_normalization(
    x: torch.Tensor,  # (B, S, T, C, D)
    valid: torch.Tensor,  # (B, S)
    irreps: torch.Tensor,  # (T, 2)
    multiplicities: torch.Tensor,  # (T,)
    running_mean: torch.Tensor,  # (T, C)
    running_var: torch.Tensor,  # (T, C)
    weight: torch.Tensor,  # (T, C)
    bias: torch.Tensor,  # (T, C)
    options: dict,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    output: (B, S, T, C, D).
    mean: (G, T, C).
    second_moment: (G, T, C).
    running_mean: (T, C).
    running_var: (T, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement equivariant_batch_normalization")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · score_model.py
• gcorso/DiffDock · all_atom_score_model.py
• gcorso/DiffDock · environment.yml
• e3nn/e3nn · _batchnorm.py
• e3nn/e3nn · _irreps.py