MolecularDiffusion.modules.layers.esen.nn.utils

Classes

BackboneInterface
HeadInterface

Functions

cg_change_mat(→ torch.tensor)
irreps_sum(→ int)	Returns the sum of the dimensions of the irreps up to the specified angular momentum.

Module Contents

class MolecularDiffusion.modules.layers.esen.nn.utils.BackboneInterface

abstractmethod forward(data: torch_geometric.data.Batch) → dict[str, torch.Tensor]

Backbone forward.

Parameters:

data (DataBatch) – Atomic systems as input

Returns:

embedding – Return backbone embeddings for the given input

Return type:

dict[str->torch.Tensor]

class MolecularDiffusion.modules.layers.esen.nn.utils.HeadInterface

abstractmethod forward(data: torch_geometric.data.Batch, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]

Head forward.

Parameters:

• data (DataBatch) – Atomic systems as input

• emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

property use_amp

MolecularDiffusion.modules.layers.esen.nn.utils.cg_change_mat(ang_mom: int, device: str = 'cpu') → torch.tensor

MolecularDiffusion.modules.layers.esen.nn.utils.irreps_sum(ang_mom: int) → int

Returns the sum of the dimensions of the irreps up to the specified angular momentum.

Parameters:

ang_mom – max angular momenttum to sum up dimensions of irreps