MolecularDiffusion.modules.models.goflow.gotennet

GotenNet: the CGR-conditioned equivariant backbone GoFlow’s velocity field is built from.

Ported from gotennet/models/representation/gotennet.py (commit 3ec00a09). Aykent & Xia, arXiv:2410.14670 – cited from the GotenNet repository’s own README, not independently verified against it (see INTEGRATION_PLAN.md, Repo Inspection).

Two deviations from upstream, both mechanical, neither changing any tensor shape or learned behaviour:

• The two einops.rearrange/reduce calls in GATA.rej (gotennet.py:296-297) are replaced by their exact plain-torch equivalents – verified by their own docstrings below.

• RankedLogger (a Hydra-rank-zero-only wrapper the rest of gotennet/utils is not ported for) is replaced by a plain logging.getLogger; the two log.info calls it backs are cosmetic (weight-init choice logging), not model behaviour.

self.distance (upstream’s Distance module) is not constructed here: see ops.py’s module docstring for why it is dead code in forward.

Attributes

log

Classes

EQFF	Equivariant feed-forward mixing. Verbatim from gotennet.py:372-414.
GATA	Attention-weighted equivariant message passing.
GotenNet	The CGR-conditioned equivariant backbone.
TimestepEmbedding	Verbatim from gotennet.py:431-447.

Functions

get_timestep_embedding(→ torch.Tensor)	Verbatim from gotennet.py:417-428.
lmax_tensor_size(→ int)
split_degree(tensor, lmax[, dim])	Split a stacked-irreps tensor into one chunk per degree l.

Module Contents

class MolecularDiffusion.modules.models.goflow.gotennet.EQFF(n_atom_basis: int, activation: Callable, epsilon: float = 1e-08, weight_init=nn.init.xavier_uniform_, bias_init=nn.init.zeros_, vec_dim=None)

Bases: torch.nn.Module

Equivariant feed-forward mixing. Verbatim from gotennet.py:372-414.

forward(s, v)

reset_parameters() → None

epsilon = 1e-08

gamma_m

n_atom_basis

w_vu

class MolecularDiffusion.modules.models.goflow.gotennet.GATA(n_atom_basis: int, activation: Callable, weight_init=nn.init.xavier_uniform_, bias_init=nn.init.zeros_, aggr='add', node_dim=0, epsilon: float = 1e-07, layer_norm='', vector_norm='', cutoff=5.0, scaling=1.0, num_heads=8, dropout=0.0, edge_updates=True, last_layer=False, scale_edge=True, edge_ln='', evec_dim=None, emlp_dim=None, sep_vecj=True, lmax=1)

Bases: torch_geometric.nn.MessagePassing

Attention-weighted equivariant message passing.

Verbatim from gotennet.py:54-370, except rej()’s two einops calls.

aggregate(features, index, ptr, dim_size)

edge_update(w1_i, w2_j, d_ij, f_ij)

forward(edge_index, s, t, dir_ij, r_ij, d_ij, num_edges_expanded)

message(edge_index, x_i, x_j, q_i, k_j, val_j, ten_j, r_ij, r_ij_attn, d_ij, dir_ij, num_edges_expanded, index, ptr: torch_geometric.typing.OptTensor, dim_size: int | None) → Tuple[torch.Tensor, torch.Tensor]

static rej(vec: torch.Tensor, d_ij: torch.Tensor) → torch.Tensor

Reject vec off the d_ij direction, per irrep degree.

Replaces gotennet.py:294-298’s two einops calls with their exact plain-torch equivalents: rearrange(d_ij, 'b l -> b l 1') is d_ij.unsqueeze(-1), and reduce(vec * d_ij_1, 'b l c -> b 1 c', 'sum') is (vec * d_ij_1).sum(dim=1, keepdim=True) – both named explicitly in INTEGRATION_PLAN.md’s Naming section.

reset_parameters() → None

update(inputs)

activation

cutoff

down_proj

dropout = 0.0

edge_mlp_dim

edge_updates = True

edge_vec_dim

epsilon = 1e-07

gamma_s

gamma_v

k_w

last_layer = False

layernorm

layernorm_ = ''

lmax = 1

n_atom_basis

num_heads = 8

phik_w_ra

q_w

scale_edge = True

sep_vecj = True

tln

update_info

vector_norm_ = ''

w_re

class MolecularDiffusion.modules.models.goflow.gotennet.GotenNet(n_atom_basis: int = 128, n_atom_feat_basis: int = 128, n_atom_rdkit_feats: int = 28, n_interactions: int = 8, radial_basis: Callable | str = 'expnorm', n_rbf: int = 20, cutoff_fn: Callable | None = None, edge_order: int = 4, activation: Callable | str | None = F.silu, max_z: int = 100, epsilon: float = 1e-08, weight_init=nn.init.xavier_uniform_, bias_init=nn.init.zeros_, max_num_neighbors: int = 32, int_layer_norm='', int_vector_norm='', num_heads=8, attn_dropout=0.0, edge_updates=True, scale_edge=True, lmax=2, aggr='add', edge_ln='', evec_dim=None, emlp_dim=None, sep_int_vec=True)

Bases: torch.nn.Module

The CGR-conditioned equivariant backbone.

Ported from gotennet.py:450-618. forward’s signature and body are unchanged from upstream: inputs reads only .edge_index, .edge_type, .batch, .r_feat, .p_feat, .atom_type off the PyG batch (never .pos – the noisy TS coordinate is always the separate x_t_N_3 argument), exactly as INTEGRATION_PLAN.md’s Repo Inspection records.

forward(x_t_N_3: torch.Tensor, t_G: torch.Tensor, inputs: Mapping[str, torch.Tensor]) → Tuple[torch.Tensor, torch.Tensor]

Compute atomic representations.

Parameters:

• x_t_N_3 – (N, 3) current (noisy) node coordinates.

• t_G – (B, 1) flow time, one value per graph.

• inputs – the PyG batch built by goflow_data.goflow_collate; only .edge_index, .edge_type, .batch, .r_feat, .p_feat, .atom_type are read.

Returns:

q is (N, hidden_dim) scalar features, mu is (N, (lmax+1)**2 - 1, hidden_dim) higher-order features.

Return type:

(q, mu)

reset_parameters() → None

atom_cgr_embedding

cutoff

cutoff_fn = None

edge_cgr_embedding

edge_embedding

edge_order = 4

eqff

gata

n_interactions = 8

neighbor_embedding

radial_basis

scale_edge = True

scaling

tensor_init

time_embedding

class MolecularDiffusion.modules.models.goflow.gotennet.TimestepEmbedding(embedding_dim, hidden_dim, output_dim)

Bases: torch.nn.Module

Verbatim from gotennet.py:431-447.

forward(timesteps: torch.Tensor) → torch.Tensor

embedding_dim

hidden_dim

mlp

output_dim

MolecularDiffusion.modules.models.goflow.gotennet.get_timestep_embedding(timesteps: torch.Tensor, embedding_dim: int, max_positions: int = 10000) → torch.Tensor

Verbatim from gotennet.py:417-428.

MolecularDiffusion.modules.models.goflow.gotennet.lmax_tensor_size(lmax: int) → int

MolecularDiffusion.modules.models.goflow.gotennet.split_degree(tensor: torch.Tensor, lmax: int, dim: int = -1)

Split a stacked-irreps tensor into one chunk per degree l.

Verbatim from gotennet.py:41-51.

MolecularDiffusion.modules.models.goflow.gotennet.log