MolecularDiffusion.modules.layers.vn.vn_layers¶

Vector Neuron (VN) equivariant building blocks.

Ported from DiffSMol (source/models/shape_vn_layers.py), which in turn derives from Deng et al., “Vector Neurons: A General Framework for SO(3)- Equivariant Networks” (ICCV 2021).

A VN feature is a tensor of shape [B, C, 3, N, ...] – C channels of 3-vectors. Every op here commutes with a global rotation applied to the 3 axis, so a stack of them is exactly SO(3)-equivariant. That is what makes the DiffSMol shape latent a (128, 3) equivariant embedding rather than an invariant descriptor.

Note these blocks are NOT translation-equivariant: get_graph_feature_cross concatenates the raw coordinate x as a channel, so the caller is responsible for centering its input first.

Classes¶

`ResnetBlockFC`	Plain (non-VN) fully connected ResNet block, used by the decoder that
`VNBatchNorm`	BatchNorm on the per-channel vector norms, leaving directions
`VNLeakyReLU`	Equivariant leaky ReLU: reflect the component of `x` that lies on
`VNLinear`	Channel-mixing linear map; acts on the channel axis only.
`VNLinearLeakyReLU`	`VNLinear` -> optional `VNBatchNorm` -> equivariant leaky ReLU.
`VNMaxPool`	Max-pool over the last axis, selecting by projection onto a learned
`VNResnetBlockFC`	Fully connected VN ResNet block.
`VNStdFeature`	Project VN features onto a learned equivariant frame, producing an

Functions¶

`get_graph_feature_cross`(→ torch.Tensor)	Build a kNN edge feature tensor from VN features.
`knn`(→ torch.Tensor)	k-nearest-neighbour indices for `x` of shape `[B, C, N]`.
`mean_pool`(→ torch.Tensor)

Module Contents¶

class MolecularDiffusion.modules.layers.vn.vn_layers.ResnetBlockFC(size_in: int, size_out: int | None = None, size_h: int | None = None)¶

Bases: torch.nn.Module

Plain (non-VN) fully connected ResNet block, used by the decoder that consumes the invariant features produced from VN latents.

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

actvn¶

fc_0¶

fc_1¶

shortcut = None¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNBatchNorm(num_features: int, dim: int)¶

Bases: torch.nn.Module

BatchNorm on the per-channel vector norms, leaving directions untouched (and therefore equivariance intact).

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

dim¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNLeakyReLU(in_channels: int, share_nonlinearity: bool = False, negative_slope: float = 0.2)¶

Bases: torch.nn.Module

Equivariant leaky ReLU: reflect the component of x that lies on the negative side of a learned direction d.

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

map_to_dir¶

negative_slope = 0.2¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNLinear(in_channels: int, out_channels: int)¶

Bases: torch.nn.Module

Channel-mixing linear map; acts on the channel axis only.

forward(x: torch.Tensor) → torch.Tensor¶: x: [B, C_in, 3, N, ...] -> [B, C_out, 3, N, ...].

map_to_feat¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNLinearLeakyReLU(in_channels: int, out_channels: int, dim: int = 5, share_nonlinearity: bool = False, negative_slope: float = 0.2, use_batchnorm: bool = True)¶

Bases: torch.nn.Module

VNLinear -> optional VNBatchNorm -> equivariant leaky ReLU.

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

dim = 5¶

map_to_dir¶

map_to_feat¶

negative_slope = 0.2¶

use_batchnorm = True¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNMaxPool(in_channels: int)¶

Bases: torch.nn.Module

Max-pool over the last axis, selecting by projection onto a learned direction (an argmax over a scalar, hence equivariant).

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

map_to_dir¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNResnetBlockFC(size_in: int, size_out: int | None = None, size_h: int | None = None)¶

Bases: torch.nn.Module

Fully connected VN ResNet block.

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

actvn_0¶

actvn_1¶

fc_0¶

fc_1¶

shortcut = None¶

class MolecularDiffusion.modules.layers.vn.vn_layers.VNStdFeature(in_channels: int, dim: int = 4, normalize_frame: bool = False, share_nonlinearity: bool = False, negative_slope: float = 0.2)¶

Bases: torch.nn.Module

Project VN features onto a learned equivariant frame, producing an invariant descriptor plus the frame itself.

forward(x: torch.Tensor)¶

dim = 4¶

normalize_frame = False¶

vn1¶

vn2¶

vn_lin¶

MolecularDiffusion.modules.layers.vn.vn_layers.get_graph_feature_cross(x: torch.Tensor, k: int = 20, idx: torch.Tensor | None = None, if_cross: bool = False) → torch.Tensor¶

Build a kNN edge feature tensor from VN features.

x: [B, C, 3, N] -> [B, 2C (or 3C), 3, N, k].

The raw x is concatenated as a channel, so this is rotation- but NOT translation-equivariant: center the point cloud before calling.

MolecularDiffusion.modules.layers.vn.vn_layers.knn(x: torch.Tensor, k: int) → torch.Tensor¶: k-nearest-neighbour indices for x of shape [B, C, N].

MolecularDiffusion.modules.layers.vn.vn_layers.mean_pool(x: torch.Tensor, dim: int = -1, keepdim: bool = False) → torch.Tensor¶