MolecularDiffusion.modules.models.ditmc.layers¶

Shared building blocks for DiTMC. Port of dit_mc/backbones/utils.py.

Parity traps carried in here from the plan, all load-bearing:

Flax nn.LayerNorm defaults to epsilon=1e-6; PyTorch’s default is 1e-5. Every LayerNorm below pins 1e-6.
Flax nn.Dense’s default kernel init is lecun_normal – a truncated normal, not PyTorch’s uniform default.
GaussianRandomFourierFeatures interleaves cos, sin (stack-then-reshape), it does not concatenate halves.
adaLN-Zero Dense layers stay zero-initialized.

Attributes¶

LAYERNORM_EPS

Classes¶

`E3MLP`	Equivariant MLP: e3x `Dense` layers with a gated activation.
`EquivariantLayerNorm`	Layer norm that respects degree/parity structure.
`GaussianRandomFourierFeatures`	`gamma(x) = [cos(2π bᵀx), sin(2π bᵀx)]` interleaved.
`MLP`	Plain (elementwise-activation) MLP. Port of `backbones/utils.MLP`.

Functions¶

`flax_dense`(→ torch.nn.Linear)	`flax.linen.Dense` equivalent: lecun-normal kernel, zero bias.
`flax_layer_norm`(→ torch.nn.LayerNorm)	`nn.LayerNorm` with Flax's defaults (eps 1e-6, last axis only).
`get_max_degree_from_tensor_e3x`(→ int)
`modulate_E3adaLN`(→ torch.Tensor)
`modulate_adaLN`(→ torch.Tensor)

Module Contents¶

class MolecularDiffusion.modules.models.ditmc.layers.E3MLP(in_features: int, num_features: int | collections.abc.Sequence[int], max_degree: int, num_parity: int, num_layers: int = 2, activation_fn: str = 'identity', use_bias: bool = True)¶

Bases: torch.nn.Module

Equivariant MLP: e3x Dense layers with a gated activation.

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

activation_fn¶

layers¶

num_layers = 2¶

out_features¶

class MolecularDiffusion.modules.models.ditmc.layers.EquivariantLayerNorm(num_features: int, max_degree: int, num_parity: int, *, use_scale: bool = True, use_bias: bool = True, epsilon: float = LAYERNORM_EPS)¶

Bases: torch.nn.Module

Layer norm that respects degree/parity structure.

The l=0 block goes through an ordinary LayerNorm. The l>0 channels are not mean-subtracted: the per-(parity, degree) norm over the order index is computed, its variance over the feature axis is taken, and the block is multiplied by rsqrt(var + eps) (times an optional learnable scales_lm). epsilon = 1e-6.

DiTMC instantiates this only with use_scale=False, use_bias=False, so in the shipped models it carries no parameters at all.

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

epsilon = 1e-06¶

has_pseudotensors¶

has_ylms¶

max_degree¶

norm00¶

num_features¶

num_parity¶

use_bias = True¶

use_scale = True¶

class MolecularDiffusion.modules.models.ditmc.layers.GaussianRandomFourierFeatures(in_features: int, features: int, sigma: float = 1.0)¶

Bases: torch.nn.Module

gamma(x) = [cos(2π bᵀx), sin(2π bᵀx)] interleaved.

b has shape (d, features//2) and is initialized normal(sigma).

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

b¶

class MolecularDiffusion.modules.models.ditmc.layers.MLP(in_features: int, num_features: int | collections.abc.Sequence[int], num_layers: int = 2, activation_fn: str = 'identity', use_bias: bool = True, output_is_zero_at_init: bool = False)¶

Bases: torch.nn.Module

Plain (elementwise-activation) MLP. Port of backbones/utils.MLP.

Uses get_activation_fn (the jax.nn.<name> family), not the e3x gated ones – DiTLayer uses this and SO3DiTLayer uses E3MLP, and they are different functions with the same names.

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

activation_fn¶

layers¶

num_layers = 2¶

out_features¶

MolecularDiffusion.modules.models.ditmc.layers.flax_dense(in_features: int, out_features: int, *, bias: bool = True, zero_init: bool = False) → torch.nn.Linear¶: flax.linen.Dense equivalent: lecun-normal kernel, zero bias.

MolecularDiffusion.modules.models.ditmc.layers.flax_layer_norm(num_features: int, *, use_scale: bool = True, use_bias: bool = True) → torch.nn.LayerNorm¶: nn.LayerNorm with Flax’s defaults (eps 1e-6, last axis only).

MolecularDiffusion.modules.models.ditmc.layers.get_max_degree_from_tensor_e3x(x: torch.Tensor) → int¶

MolecularDiffusion.modules.models.ditmc.layers.modulate_E3adaLN(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) → torch.Tensor¶

MolecularDiffusion.modules.models.ditmc.layers.modulate_adaLN(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) → torch.Tensor¶

MolecularDiffusion.modules.models.ditmc.layers.LAYERNORM_EPS = 1e-06¶