MolecularDiffusion.modules.layers.tabasco.positional_encoder

Classes

PositionalEncoding	Abstract interface for modules that add positional information to tensors.
SinusoidEncoding	Classic sinusoidal positional encoding (Vaswani et al., 2017).
TimeFourierEncoding	Encoder for continuous timesteps in [0, 1]

Module Contents

class MolecularDiffusion.modules.layers.tabasco.positional_encoder.PositionalEncoding

Bases: abc.ABC, torch.nn.Module

Abstract interface for modules that add positional information to tensors.

abstractmethod forward(*args, **kwargs) → torch.Tensor

Return positional encodings.

abstractmethod out_dim()

Embedding dimension produced by this encoder.

class MolecularDiffusion.modules.layers.tabasco.positional_encoder.SinusoidEncoding(posenc_dim, max_len=100, random_permute=False)

Bases: PositionalEncoding

Classic sinusoidal positional encoding (Vaswani et al., 2017).

Initialize the positional encoding.

Parameters:

• posenc_dim – Size of the embedding dimension.

• max_len – Maximum sequence length supported. seq_len passed to forward must not exceed this value.

• random_permute – If True, the positions are randomly permuted for each sample (useful as lightweight data augmentation but destroys absolute ordering).

forward(batch_size: int, seq_len: int)

Return positional embeddings of shape (batch_size, seq_len, posenc_dim).

Parameters:

• batch_size – Number of samples in the batch.

• seq_len – Length of the sequence.

Note

seq_len must not exceed the max_len passed at construction.

out_dim()

Embedding dimension produced by this encoder.

max_len = 100

posenc_dim

random_permute = False

class MolecularDiffusion.modules.layers.tabasco.positional_encoder.TimeFourierEncoding(posenc_dim, max_len=100, random_permute=False)

Bases: PositionalEncoding

Encoder for continuous timesteps in [0, 1]

forward(t: torch.Tensor)

Encode a tensor of timesteps.

Parameters:

t – 1-D tensor with values in [0, 1].

Returns:

Tensor of shape (B, posenc_dim) with sine/cosine features.

out_dim()

Embedding dimension produced by this encoder.

max_len = 100

posenc_dim

random_permute = False