MolecularDiffusion.modules.models.equiformer_v2_backbone

EquiformerV2Backbone — plain per-node scalar-feature backbone for property prediction. Unlike EquiformerV2_dynamics (built for denoising: mandatory timestep, velocity + feature-delta output), this exposes per-node l=0 features via forward(data) -> {“x”: …}, mirroring eSEN_Backbone’s contract so it plugs into ProperyPrediction.

The EquiformerV2 encoder and all layer files are NOT modified.

Classes

EquiformerV2Backbone

Module Contents

class MolecularDiffusion.modules.models.equiformer_v2_backbone.EquiformerV2Backbone(equiformer: torch.nn.Module, in_node_channels: int, sphere_channels: int = 128, lmax_list=None)

Bases: torch.nn.Module

Parameters:

  • equiformer – Instantiated EquiformerV2 model (from shepherd_arch).

  • in_node_channels (int) – Width of the per-node input feature, i.e. atomic_number(1) + data.x width.

  • sphere_channels (int) – Must match equiformer.sphere_channels.

  • lmax_list (list[int]) – Must match equiformer.lmax_list.

forward(data) → dict[str, torch.Tensor]
d_model = 128
equiformer
input_proj
lmax_list = [6]
sphere_channels = 128