MolecularDiffusion.modules.models.goflow.outputs

The atom-wise 3D output head that turns (q, mu) into a velocity.

Ported from gotennet/models/components/outputs.py:21-49 (SNNDense, what GatedEquivariantBlock needs) and :52-189 (GatedEquivariantBlock, Atomwise3DOut). Atomwise3DOutRTSP (upstream’s own “TODO: modify to accept r/ts/p”, never imported by flow_matching/flow_module.py) is not ported – dead code, not a fidelity gap.

Classes

Atomwise3DOut	Two stacked GatedEquivariantBlocks mapping ``(q, mu[:,
GatedEquivariantBlock	Rotationally invariant/equivariant tensorial feature mixing.
SNNDense	Fully connected linear layer with activation. Verbatim from

Module Contents

class MolecularDiffusion.modules.models.goflow.outputs.Atomwise3DOut(n_in, n_hidden: int | None = None, activation=shifted_softplus)

Bases: torch.nn.Module

Two stacked GatedEquivariantBlocks mapping (q, mu[:, :3, :]) to a per-atom 3D velocity. Verbatim from outputs.py:161-189.

forward(l0: torch.Tensor, l1: torch.Tensor) → torch.Tensor

out_net

class MolecularDiffusion.modules.models.goflow.outputs.GatedEquivariantBlock(n_sin, n_vin, n_sout, n_vout, n_hidden, activation=F.silu, sactivation=None)

Bases: torch.nn.Module

Rotationally invariant/equivariant tensorial feature mixing.

Verbatim from outputs.py:52-85.

forward(scalars: torch.Tensor, vectors: torch.Tensor)

mix_vectors

n_hidden

n_sin

n_sout

n_vin

n_vout

sactivation = None

scalar_net

class MolecularDiffusion.modules.models.goflow.outputs.SNNDense(in_features: int, out_features: int, bias: bool = True, activation: Callable | torch.nn.Module | None = None, weight_init: Callable = xavier_uniform_, bias_init: Callable = zeros_)

Bases: torch.nn.Linear

Fully connected linear layer with activation. Verbatim from outputs.py:21-49.

forward(input: torch.Tensor) → torch.Tensor

reset_parameters() → None

activation

bias_init

weight_init