MolecularDiffusion.modules.models.ditmc.readout

Readouts. Port of dit_mc/backbones/readout.py.

Note a quirk that is reproduced deliberately: neither readout applies cond_scaling. The conditioner’s features enter the readout unscaled, so classifier-free guidance does not zero them here even when it zeroes them in every DiT block. That is upstream behaviour.

Neither readout branches on act_dense_correct_bool either – both always use act_fn(Dense(c)).

Classes

EquivariantReadout	SO(3)-equivariant readout (dit_so3).
SimpleReadout	Non-equivariant readout (dit_ape / dit_rpe).

Module Contents

class MolecularDiffusion.modules.models.ditmc.readout.EquivariantReadout(num_features: int, activation_fn: str, output: str = 'drift_and_noise', in_max_degree: int = 1, in_num_parity: int = 2)

Bases: torch.nn.Module

SO(3)-equivariant readout (dit_so3).

y[:, 0] + y[:, 1] deliberately mixes the even and odd parity blocks – it breaks O(3) but leaves the output SO(3)-equivariant, which is what a coordinate model needs. The l=1 block (indices 1:4) is the answer; index 0 (the scalar) is discarded.

forward(features_nodes: torch.Tensor, features_time: torch.Tensor, features_cond: torch.Tensor | None = None)

act_fn

ada_dense

cond_norm

head

in_max_degree = 1

in_num_parity = 2

norm

num_features

output = 'drift_and_noise'

class MolecularDiffusion.modules.models.ditmc.readout.SimpleReadout(num_features: int, activation_fn: str, output: str = 'drift_and_noise')

Bases: torch.nn.Module

Non-equivariant readout (dit_ape / dit_rpe).

forward(features_nodes: torch.Tensor, features_time: torch.Tensor, features_cond: torch.Tensor | None = None)

act_fn

ada_dense

head

norm

num_features

output = 'drift_and_noise'