MolecularDiffusion.modules.tasks.diffusion_tabasco_equiformer

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

Replaces TABASCO’s TransformerModule backbone with EquiformerV2TabascoBackbone (modules/models/tabasco_equiformer/equiformer_backbone.py, wrapping this platform’s existing EquiformerV2 encoder). 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_equiformer/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 – EquiformerV2TabascoBackbone(**equiformer_config) in place of TransformerModule(**transformer_config) – the exact pattern TabascoEGNNDiffusionTask (diffusion_tabasco_egnn.py:50-97) already establishes.

Classes

ModelTaskFactory	Factory for TabascoEquiformerDiffusionTask.
TabascoEquiformerDiffusionTask	TABASCO flow-matching diffusion with an EquiformerV2 backbone.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_tabasco_equiformer.ModelTaskFactory(task_type: str, equiformer_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 TabascoEquiformerDiffusionTask.

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 (equiformer_config in place of transformer_config) and build()’s target class differ.

build()

Build and return the TabascoEquiformerDiffusionTask.

atom_vocab

atomics_interpolant_config

coords_interpolant_config

dataset_stats

equiformer_config

flow_matching_config

kwargs

num_atom_types

task_type

train_set = None

class MolecularDiffusion.modules.tasks.diffusion_tabasco_equiformer.TabascoEquiformerDiffusionTask(equiformer_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 EquiformerV2 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_equiformer'

to_pointcloud

to_tensordict