MolecularDiffusion.modules.models.diffsmol.common¶

Shared primitives for the DiffSMol backbone.

Ported from DiffSMol source/models/common.py (only the parts the bond-free score model actually reaches).

GVP/GVPLayerNorm here are DiffSMol’s own variants and are NOT interchangeable with modules/layers/gvp (FlowMol’s): that one takes a four-dim (dim_vectors_in, dim_vectors_out, dim_feats_in, dim_feats_out) signature with SiLU/Sigmoid activations and a different tensor layout, while this one takes (in_dims, h_dim, out_dims) tuples with ReLU/sigmoid. They are kept separate rather than unified because silently swapping activations inside a ported architecture is how you get a model that trains but is not the model you ported.

Classes¶

`GVP`	Geometric Vector Perceptron (Jing et al. 2021), vector-gated.
`GVPLayerNorm`	LayerNorm on scalars; RMS rescale (no learned params) on vectors.
`GaussianSmearing`	Distance expansion on a fixed, hand-chosen offset grid.
`MLP`	MLP with the same hidden dim across all layers.
`ShiftedSoftplus`
`SinusoidalPosEmb`	Sinusoidal timestep embedding.

Functions¶

`norm_no_nan`(→ torch.Tensor)	L2 norm clamped below at `eps` so the gradient never sees 0/0.
`outer_product`(→ torch.Tensor)	Flattened outer product of a sequence of per-edge feature tensors.

Module Contents¶

class MolecularDiffusion.modules.models.diffsmol.common.GVP(in_dims: tuple[int, int], h_dim: int, out_dims: tuple[int, int], activations=(F.relu, torch.sigmoid), vector_gate: bool = True)¶

Bases: torch.nn.Module

Geometric Vector Perceptron (Jing et al. 2021), vector-gated.

Takes and returns (scalars [N, s], vectors [N, v, 3]).

forward(x)¶

dummy_param¶

vector_gate = True¶

class MolecularDiffusion.modules.models.diffsmol.common.GVPLayerNorm(dims: tuple[int, int])¶

Bases: torch.nn.Module

LayerNorm on scalars; RMS rescale (no learned params) on vectors.

forward(x)¶

scalar_norm¶

class MolecularDiffusion.modules.models.diffsmol.common.GaussianSmearing(start: float = 0.0, stop: float = 5.0, num_gaussians: int = 50)¶

Bases: torch.nn.Module

Distance expansion on a fixed, hand-chosen offset grid.

Note the offsets are hardcoded upstream (20 of them), so num_gaussians is advisory only – the true feature width is always 20. It is kept in the signature because the config sets num_r_gaussian: 20 to match.

forward(dist: torch.Tensor) → torch.Tensor¶

coeff¶

num_gaussians¶

start = 0.0¶

stop = 5.0¶

class MolecularDiffusion.modules.models.diffsmol.common.MLP(in_dim: int, out_dim: int, hidden_dim: int, num_layer: int = 2, norm: bool = True, act_fn: str = 'relu', act_last: bool = False)¶

Bases: torch.nn.Module

MLP with the same hidden dim across all layers.

forward(x: torch.Tensor) → torch.Tensor¶

net¶

class MolecularDiffusion.modules.models.diffsmol.common.ShiftedSoftplus¶

Bases: torch.nn.Module

forward(x: torch.Tensor) → torch.Tensor¶

shift¶

class MolecularDiffusion.modules.models.diffsmol.common.SinusoidalPosEmb(dim: int)¶

Bases: torch.nn.Module

Sinusoidal timestep embedding.

forward(x: torch.Tensor) → torch.Tensor¶

dim¶

MolecularDiffusion.modules.models.diffsmol.common.norm_no_nan(x: torch.Tensor, axis: int = -1, keepdims: bool = False, eps: float = 1e-08, sqrt: bool = True) → torch.Tensor¶: L2 norm clamped below at eps so the gradient never sees 0/0.

MolecularDiffusion.modules.models.diffsmol.common.outer_product(*vectors: torch.Tensor) → torch.Tensor¶: Flattened outer product of a sequence of per-edge feature tensors.