MolecularDiffusion.modules.tasks.diffusion_tabasco_egnn

TABASCO + EGNN backbone: a novel-model derivation of diffusion_tabasco.

Replaces TABASCO’s TransformerModule backbone with EGNNBackbone (modules/models/tabasco_egnn/egnn_backbone.py, wrapping this platform’s existing EGNN/GCL/EquivariantUpdate/EquivariantBlock layers). Every other TABASCO component – FlowMatchingModel, SDEMetricInterpolant, DiscreteInterpolant, the Euler-Maruyama sampler, TabascoNodeDistribution, the pointcloud<->TensorDict adapters – is imported unmodified from MolecularDiffusion.modules.tasks.diffusion_tabasco.

See docs/model_novel/tabasco_egnn/INTEGRATION_PLAN.md (“Integration Plan”, “Derivation Rung”) for why this is a subclass that overrides only __init__ rather than a smaller override: TabascoDiffusionTask.__init__ has no seam to override just the backbone-construction line, so this reproduces its assembly sequence (diffusion_tabasco.py:280-326) with one substitution – EGNNBackbone(**egnn_config) in place of TransformerModule(**transformer_config).

Classes

ModelTaskFactory	Factory for TabascoEGNNDiffusionTask.
TabascoEGNNDiffusionTask	TABASCO flow-matching diffusion with an EGNN backbone.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_tabasco_egnn.ModelTaskFactory(task_type: str, egnn_config: dict, coords_interpolant_config: dict, atomics_interpolant_config: dict, flow_matching_config: dict, num_atom_types: int, dataset_stats: dict, atom_vocab: list | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs)

Bases: MolecularDiffusion.modules.tasks.diffusion_tabasco.ModelTaskFactory

Factory for TabascoEGNNDiffusionTask.

compute_dataset_stats is inherited unchanged from TabascoModelTaskFactory (diffusion_tabasco.py:191-245) – it only touches self.dataset_stats/self.train_set, neither of which changes shape here. Only the constructor’s config surface (egnn_config in place of transformer_config) and build()’s target class differ.

build()

Build and return the TabascoEGNNDiffusionTask.

atom_vocab

atomics_interpolant_config

coords_interpolant_config

dataset_stats

egnn_config

flow_matching_config

kwargs

num_atom_types

task_type

train_set = None

class MolecularDiffusion.modules.tasks.diffusion_tabasco_egnn.TabascoEGNNDiffusionTask(egnn_config: dict, coords_interpolant_config: dict, atomics_interpolant_config: dict, flow_matching_config: dict, num_atom_types: int, dataset_stats: dict, atom_vocab: list | None = None, condition_names: list = [], context_mask_rate: float = 0.0, mask_value: float = 0.0, normalize_condition: str | None = None, adapter_conditions: list | None = None, use_adapter_module: bool = False)

Bases: MolecularDiffusion.modules.tasks.diffusion_tabasco.TabascoDiffusionTask

TABASCO flow-matching diffusion with an EGNN backbone.

Every method other than __init__ (forward, predict_and_target, evaluate, sample, node_dist_model, n_node_dist, model, device) is inherited unchanged from TabascoDiffusionTask – they only ever call through self.tabasco_model/self.net, which is backbone-agnostic.

atom_vocab = None

condition = []

context_mask_rate = 0.0

mask_value = 0.0

max_n_nodes

n_adapter_context

n_concat_context

normalize_condition = None

num_atom_types

prop_dist_model = None

property_norms = None

tabasco_model

task_type = 'diffusion_tabasco_egnn'

to_pointcloud

to_tensordict