MolecularDiffusion.modules.models.gfmdiff.equignn¶
GFMDiff Dual-Track Transformer network body.
Near-verbatim port of GFMDiff/models/diff/dtn.py::EquiGNN. Only the
import paths changed (blocks now live under
MolecularDiffusion.modules.layers.gfmdiff.blocks); the math/forward
logic is unchanged from the source repo.
Per the approved integration plan, the diffused “degree” (valence) channel
and the Geometric-Facilitated Loss are out of scope for this integration:
include_de is hardcoded to False below (the constructor no longer
reads it from model_config), while include_an (atomic-number
channel) is kept, matching every other backbone already integrated in this
repo (in_node_nf = len(atom_vocab) + n_dim_extra + 1).
Classes¶
Module Contents¶
- class MolecularDiffusion.modules.models.gfmdiff.equignn.EquiGNN(model_config, data_config)¶
Bases:
torch.nn.Module- abstractmethod forward(z_t, t, node_mask, pair_mask, context=None)¶
- unwrap_forward()¶
- wrap_forward(node_mask, edge_mask, context)¶
- act¶
- activate¶
- add_time¶
- angle_emb¶
- block_calc¶
- context_dim¶
- convs¶
- dropout¶
- edge_emb¶
- emb_dim¶
- include_an¶
- include_de = False¶
- node_nf¶
- node_pair_fuse¶
- num_heads¶
- num_layers¶
- pair_scale¶
- xh_embedding¶
- xh_embedding_out¶