MolecularDiffusion.modules.models.chefnmr.layers¶

Upstream: src/model/modules/layers.py, modified from DiT (facebookresearch/DiT) and timm’s VisionTransformer.

Kept in modules/models/chefnmr/ rather than promoted to modules/layers/: these blocks carry an explicit atom mask through attention and use adaLN-zero conditioning on a fused time + spectrum vector. Nothing else in the tree consumes that shape today, and a shared block would have to grow the mask argument anyway.

Attributes¶

to_2tuple

Classes¶

`Attention`	Masked multi-head self-attention over the atom axis.
`DiTBlock`	DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
`FinalLayer`	adaLN-zero output head: hidden -> coordinate update.
`Mlp`	timm's ViT MLP.
`TimestepEmbedder`	Sinusoidal noise-level embedding + 2-layer MLP.

Functions¶

modulate(→ torch.Tensor)

Module Contents¶

class MolecularDiffusion.modules.models.chefnmr.layers.Attention(dim: int, num_heads: int = 8, qkv_bias: bool = False, qk_norm: bool = False, proj_bias: bool = True, attn_drop: float = 0.0, proj_drop: float = 0.0, norm_layer: Type[torch.nn.Module] = nn.LayerNorm)¶

Bases: torch.nn.Module

Masked multi-head self-attention over the atom axis.

fused_attn is hard-off, exactly as upstream: the released weights were trained through the explicit softmax path, and scaled_dot_product_attention differs in how it handles a fully masked row.

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

attn_drop¶

fused_attn = False¶

head_dim¶

k_norm¶

num_heads = 8¶

proj¶

proj_drop¶

q_norm¶

qkv¶

scale¶

class MolecularDiffusion.modules.models.chefnmr.layers.DiTBlock(hidden_size: int, n_heads: int, mlp_ratio: float = 4, **block_kwargs)¶

Bases: torch.nn.Module

DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.

forward(x: torch.Tensor, c: torch.Tensor, atom_mask: torch.Tensor) → torch.Tensor¶

adaLN_modulation¶

attn¶

mlp¶

norm1¶

norm2¶

class MolecularDiffusion.modules.models.chefnmr.layers.FinalLayer(hidden_size: int, out_atom_coords_size: int)¶

Bases: torch.nn.Module

adaLN-zero output head: hidden -> coordinate update.

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

adaLN_modulation¶

linear¶

norm_final¶

class MolecularDiffusion.modules.models.chefnmr.layers.Mlp(in_features: int, hidden_features: int = None, out_features: int = None, act_layer=nn.GELU, norm_layer=None, bias: bool = True, drop: float = 0.0, use_conv: bool = False)¶