MolecularDiffusion.modules.layers.tabasco.transformer_module

Classes

TransformerModule

Basic Transformer model for molecule generation.

Module Contents

class MolecularDiffusion.modules.layers.tabasco.transformer_module.TransformerModule(spatial_dim: int, atom_dim: int, num_heads: int, num_layers: int, hidden_dim: int, activation: str = 'SiLU', implementation: str = 'pytorch', cross_attention: bool = False, add_sinusoid_posenc: bool = True, concat_combine_input: bool = False, custom_weight_init: str | None = None)

Bases: torch.nn.Module

Basic Transformer model for molecule generation.

Parameters:

• custom_weight_init – None, “xavier”, “kaiming”, “orthogonal”, “uniform”, “eye”, “normal”

• well) ((uniform does not work)

forward(coords, atomics, padding_mask, t) → torch.Tensor

Forward pass of the module.

add_sinusoid_posenc = True

atom_type_embed

comb_input_dim

concat_combine_input = False

cross_attention = False

custom_weight_init = None

hidden_dim

implementation = 'pytorch'

input_dim

linear_embed

num_layers

out_atom_type_linear

out_coord_linear

time_dim = 1

time_encoding