MolecularDiffusion.modules.models.ditmc.embedding¶

Node, edge and time embeddings. Port of dit_mc/backbones/embedding.py.

Attributes¶

`NUM_ATOMIC_NUMBERS`
`NUM_SHORTEST_HOPS`
`UNREACHABLE_HOPS`

Classes¶

`DiTEdgeEmbed`	Non-equivariant relative edge embedding (`dit_ape` / `dit_rpe`).
`DiTNodeEmbed`	Atomic-number embedding, optionally plus self-conditioning / positions.
`EdgeAttributeEmbedding`	MeshGraphNet edge encoder over the 4-class bond one-hot.
`NodeAttributeEmbedding`	MeshGraphNet node encoder: 2-layer MLP + LayerNorm.
`RadialSphericalEdgeEmbedding`	SO(3)-equivariant edge basis (`dit_so3`).
`TimeEmbedding`	Fourier features of the latent time, then a 2-layer MLP.

Functions¶

`build_cond_embeddings`(→ tuple[NodeAttributeEmbedding, ...)	Convenience pair used by the MeshGraphNet encoder.
`flax_embedding`(→ torch.nn.Embedding)	`flax.linen.Embed`: `variance_scaling(1.0, 'fan_in', 'normal')`.
`get_index_embedding`(→ torch.Tensor)	Sinusoidal positional embedding from per-molecule atom offsets.
`get_pos_indices`(→ torch.Tensor)	Index of each atom within its own molecule.

Module Contents¶

class MolecularDiffusion.modules.models.ditmc.embedding.DiTEdgeEmbed(num_features: int, activation_fn: str, embed_distances_bool: bool = True, embed_shortest_hops_bool: bool = True, radial_basis_bool: bool | None = None, num_radial_basis: int | None = None, max_frequency: float | None = None)¶

Bases: torch.nn.Module

Non-equivariant relative edge embedding (dit_ape / dit_rpe).

embed_shortest_hops_bool is True in globals, which is why relative_embedding_bool is True for both dit_ape and dit_rpe; they differ only in embed_distances_bool.

The shortest-hops MLP deliberately keeps use_bias=True – upstream’s own comment: “we need a bias here s.t. the output is non-zero in case of CFG”.

forward(graph: MolecularDiffusion.modules.models.ditmc.graphs.LatentGraph) → torch.Tensor¶

embed_distances_bool = True¶

embed_shortest_hops_bool = True¶

max_frequency = None¶

num_features¶

num_radial_basis = None¶

radial_basis_bool = False¶

class MolecularDiffusion.modules.models.ditmc.embedding.DiTNodeEmbed(num_features: int, activation_fn: str, self_conditioning_bool: bool, positional_encoding_bool: bool, positional_embedding_bool: bool)¶

Bases: torch.nn.Module

Atomic-number embedding, optionally plus self-conditioning / positions.

forward(graph: MolecularDiffusion.modules.models.ditmc.graphs.LatentGraph) → torch.Tensor¶

embed¶

num_features¶

pos_mlp¶

positional_embedding_bool¶

positional_encoding_bool¶

sc_mlp¶

self_conditioning_bool¶

class MolecularDiffusion.modules.models.ditmc.embedding.EdgeAttributeEmbedding(in_features: int, num_features: int, activation_fn: str)¶

Bases: torch.nn.Module

MeshGraphNet edge encoder over the 4-class bond one-hot.

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

mlp¶

norm¶

class MolecularDiffusion.modules.models.ditmc.embedding.NodeAttributeEmbedding(in_features: int, num_features: int, activation_fn: str)¶

Bases: torch.nn.Module

MeshGraphNet node encoder: 2-layer MLP + LayerNorm.

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

mlp¶

norm¶

class MolecularDiffusion.modules.models.ditmc.embedding.RadialSphericalEdgeEmbedding(max_degree: int, activation_fn: str, cutoff: float | None = None, embed_shortest_hops_bool: bool = False, scale_spherical_basis_with_shortest_hops_bool: bool = False, radial_basis: str = 'reciprocal_bernstein', num_radial_basis: int = 32, radial_basis_kwargs: dict | None = None, cutoff_fn: str | None = None)¶

Bases: torch.nn.Module

SO(3)-equivariant edge basis (dit_so3).

e3x.nn.basis with reciprocal_bernstein radial functions and no cutoff (dit_so3 sets cutoff: null; the factory raises otherwise). The shortest-hops term is zero-initialized so it starts as a pure identity scaling, and is gated by cond_scaling for classifier-free guidance.

forward(graph: MolecularDiffusion.modules.models.ditmc.graphs.LatentGraph) → torch.Tensor¶

embed_shortest_hops_bool = False¶

max_degree¶

num_radial_basis = 32¶

radial_fn¶

scale_with_hops = False¶

class MolecularDiffusion.modules.models.ditmc.embedding.TimeEmbedding(num_features: int, num_features_fourier: int | None = None, activation_fn: str = 'silu')¶

Bases: torch.nn.Module

Fourier features of the latent time, then a 2-layer MLP.

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

ff¶

mlp¶

MolecularDiffusion.modules.models.ditmc.embedding.build_cond_embeddings(node_attr_dim: int, num_features: int, activation_fn: str) → tuple[NodeAttributeEmbedding, EdgeAttributeEmbedding]¶: Convenience pair used by the MeshGraphNet encoder.

MolecularDiffusion.modules.models.ditmc.embedding.flax_embedding(num_embeddings: int, features: int) → torch.nn.Embedding¶

flax.linen.Embed: variance_scaling(1.0, 'fan_in', 'normal').

For a (num_embeddings, features) table with out_axis=0 this resolves to fan_in = features and a plain (not truncated) normal.

MolecularDiffusion.modules.models.ditmc.embedding.get_index_embedding(indices: torch.Tensor, emb_dim: int, max_len: int = 256) → torch.Tensor¶: Sinusoidal positional embedding from per-molecule atom offsets.

MolecularDiffusion.modules.models.ditmc.embedding.get_pos_indices(batch_segments: torch.Tensor, num_graphs: int) → torch.Tensor¶: Index of each atom within its own molecule.

MolecularDiffusion.modules.models.ditmc.embedding.NUM_ATOMIC_NUMBERS = 119¶

MolecularDiffusion.modules.models.ditmc.embedding.NUM_SHORTEST_HOPS = 512¶

MolecularDiffusion.modules.models.ditmc.embedding.UNREACHABLE_HOPS = 510¶