MolecularDiffusion.modules.models.midi.placeholder

MiDi’s multi-modality tensor container.

Trimmed from upstream midi/utils.py: the wandb setup, the torchmetrics NoSync* wrappers and to_dense (the PyG -> dense bridge) are all gone. to_dense in particular is replaced wholesale by the platform’s graph3d_dense_collate, which already emits these exact shapes.

Classes

Dims	Per-modality channel counts.
PlaceHolder	Dense batch of the four diffused modalities plus the global feature.

Functions

remove_mean_with_mask(→ torch.Tensor)

Subtract the per-molecule mean over real nodes only.

Module Contents

class MolecularDiffusion.modules.models.midi.placeholder.Dims

Per-modality channel counts.

Upstream reuses PlaceHolder itself as the dims container, which reads badly (a pos field holding the integer 3). A dataclass says the same thing without pretending to be a batch.

E: int

X: int

charges: int

pos: int = 3

y: int

class MolecularDiffusion.modules.models.midi.placeholder.PlaceHolder(pos: torch.Tensor | None, X: torch.Tensor | None, charges: torch.Tensor | None, E: torch.Tensor | None, y: torch.Tensor | None, t_int: torch.Tensor | None = None, t: torch.Tensor | None = None, node_mask: torch.Tensor | None = None)

Dense batch of the four diffused modalities plus the global feature.

pos (B,N,3), X (B,N,K), charges (B,N,C), E (B,N,N,5), y (B,dy), node_mask (B,N) bool.

collapse(collapse_charges: torch.Tensor) → PlaceHolder

One-hot/logit modalities -> integer class ids (charges decoded).

Padded entries are marked out of range exactly as upstream does (X = -1, charges = 1000, E = -1); the task clamps them before handing anything to the platform.

copy() → PlaceHolder

Shallow copy sharing the underlying tensors.

device_as(x: torch.Tensor) → PlaceHolder

Move every present modality onto x’s device.

mask(node_mask: torch.Tensor | None = None) → PlaceHolder

Zero padded rows/columns and the diagonal, and re-centre pos.

The trailing symmetry check is upstream’s and is kept: MiDi’s whole edge pipeline assumes E == E^T, and a violation shows up as bad molecules rather than as an exception anywhere else.

E

X

charges

node_mask = None

pos

t = None

t_int = None

y

MolecularDiffusion.modules.models.midi.placeholder.remove_mean_with_mask(x: torch.Tensor, node_mask: torch.Tensor) → torch.Tensor

Subtract the per-molecule mean over real nodes only.

Parameters:

• x – (B, N, D) tensor, already zeroed on padded rows.

• node_mask – (B, N) bool, True for real atoms.

Returns:

x with the masked mean removed from every real row.