base

base

PositionSource

Bases: Enum

Where positional encodings read positions from.

DenominatorMode

Bases: Enum

Frequency denominator convention for sinusoidal encodings.

OrderingMode

Bases: Enum

Sine/cosine channel ordering convention.

PositionalEncoding

PositionalEncoding(embedding_dimension, precompute_encodings=True, maximum_sequence_length=5000, mlp_hidden_dimensions=None, mlp_activation=SiLU)

Bases: ABC, Module

Base class for positional encoding with optional precomputing and MLP learnable layer.

Source code in src/versatil/models/layers/positional_encoding/base.py

def __init__(
    self,
    embedding_dimension: int,
    precompute_encodings: bool = True,
    maximum_sequence_length: int | None = 5000,
    mlp_hidden_dimensions: list[int] | None = None,
    mlp_activation: Callable | None = nn.SiLU,
):
    super().__init__()
    self.embedding_dimension = embedding_dimension
    self.maximum_sequence_length = maximum_sequence_length
    self.precompute_encodings = precompute_encodings
    self.mlp_network = (
        None  # An extra learnable MLP layer after positional encoding.
    )
    if mlp_hidden_dimensions:
        activation = mlp_activation if mlp_activation is not None else nn.SiLU
        self.mlp_network = MLP(
            input_dimension=embedding_dimension,
            hidden_dimensions=mlp_hidden_dimensions[:-1],
            output_dim=mlp_hidden_dimensions[-1],
            activation_function=activation,
        )

forward abstractmethod

forward(input_tensor, offset=0)

Compute positional encodings for the input tensor.

Source code in src/versatil/models/layers/positional_encoding/base.py

@abstractmethod
def forward(self, input_tensor: torch.Tensor, offset: int = 0) -> torch.Tensor:
    """Compute positional encodings for the input tensor."""
    raise NotImplementedError("Subclasses must implement forward")

PositionalEncoding1D

PositionalEncoding1D(embedding_dimension, position_source=TENSOR_INDICES.value, precompute_encodings=True, maximum_sequence_length=5000, mlp_hidden_dimensions=None, mlp_activation=SiLU)

Bases: PositionalEncoding, ABC

Base class for 1D positional encodings.

Source code in src/versatil/models/layers/positional_encoding/base.py

def __init__(
    self,
    embedding_dimension: int,
    position_source: str = PositionSource.TENSOR_INDICES.value,
    precompute_encodings: bool = True,
    maximum_sequence_length: int | None = 5000,
    mlp_hidden_dimensions: list[int] | None = None,
    mlp_activation: Callable | None = nn.SiLU,
):
    self.position_source = position_source
    super().__init__(
        embedding_dimension=embedding_dimension,
        precompute_encodings=precompute_encodings,
        maximum_sequence_length=maximum_sequence_length,
        mlp_hidden_dimensions=mlp_hidden_dimensions,
        mlp_activation=mlp_activation,
    )
    if (
        precompute_encodings
        and self.position_source == PositionSource.TENSOR_INDICES.value
    ):
        if self.maximum_sequence_length is None:
            raise ValueError(
                "maximum_sequence_length must be set when precompute_encodings=True"
            )
        precomputed_encodings = self._compute_encodings(
            torch.arange(self.maximum_sequence_length).float()
        )
        self.register_buffer(
            "precomputed_encodings", precomputed_encodings.unsqueeze(0)
        )  # [1, maximum_sequence_length, embedding_dimension]

forward

forward(input_tensor, offset=0)

Compute positional encodings for input tensor.

Parameters:

Name	Type	Description	Default
input_tensor	Tensor	Input tensor with batch-first convention. - For TENSOR_INDICES: shape (batch_size, seq_len, ...) or (batch_size, seq_len) - For SCALAR: shape (batch_size,) containing scalar values to encode	required
offset	int	Position offset for TENSOR_INDICES mode. Shifts indices from [0..seq_len-1] to [offset..offset+seq_len-1]. Used during cached autoregressive inference.	0

Returns:

Type	Description
Tensor	Positional encodings with shape: - For TENSOR_INDICES: (batch_size, seq_len, embedding_dimension) - For SCALAR: (batch_size, embedding_dimension)

Source code in src/versatil/models/layers/positional_encoding/base.py

def forward(self, input_tensor: torch.Tensor, offset: int = 0) -> torch.Tensor:
    """Compute positional encodings for input tensor.

    Args:
        input_tensor: Input tensor with batch-first convention.
            - For TENSOR_INDICES: shape (batch_size, seq_len, ...) or (batch_size, seq_len)
            - For SCALAR: shape (batch_size,) containing scalar values to encode
        offset: Position offset for TENSOR_INDICES mode. Shifts indices from
            [0..seq_len-1] to [offset..offset+seq_len-1]. Used during cached
            autoregressive inference.

    Returns:
        Positional encodings with shape:
            - For TENSOR_INDICES: (batch_size, seq_len, embedding_dimension)
            - For SCALAR: (batch_size, embedding_dimension)
    """
    encodings: torch.Tensor
    if self.position_source == PositionSource.TENSOR_INDICES.value:
        batch_size = input_tensor.size(0)
        seq_len = input_tensor.size(1)
        if self.precompute_encodings:
            if (
                self.maximum_sequence_length is not None
                and offset + seq_len > self.maximum_sequence_length
            ):
                raise ValueError(
                    f"Requested positions [{offset}, {offset + seq_len}) exceed "
                    f"precomputed maximum_sequence_length {self.maximum_sequence_length}. "
                    f"Increase maximum_sequence_length."
                )
            encodings = self.precomputed_encodings[
                :, offset : offset + seq_len, :
            ]  # [1, seq_len, embedding_dimension]
        else:
            encodings = self._compute_encodings(
                torch.arange(offset, offset + seq_len).to(input_tensor.device)
            )
            encodings = encodings.unsqueeze(0)  # [1, seq_len, embedding_dimension]
        encodings = encodings.expand(
            batch_size, -1, -1
        )  # [batch_size, seq_len, embedding_dimension]
    elif self.position_source == PositionSource.SCALAR.value:
        encodings = self._compute_encodings(
            input_tensor
        )  # [batch_size, embedding_dimension]
    else:
        raise ValueError(
            f"Unsupported position_source for 1D: {self.position_source}"
        )
    if self.mlp_network:
        encodings_mlp: torch.Tensor = self.mlp_network(encodings)
        return encodings_mlp
    return encodings

PositionalEncoding2D

PositionalEncoding2D(embedding_dimension, mlp_hidden_dimensions=None, mlp_activation=SiLU)

Bases: PositionalEncoding, ABC

Base class for 2D positional encodings.

Source code in src/versatil/models/layers/positional_encoding/base.py

def __init__(
    self,
    embedding_dimension: int,
    mlp_hidden_dimensions: list[int] | None = None,
    mlp_activation: Callable | None = nn.SiLU,
):
    super().__init__(
        embedding_dimension=embedding_dimension,
        precompute_encodings=False,  # No precompute for variable 2D shapes
        mlp_hidden_dimensions=mlp_hidden_dimensions,
        mlp_activation=mlp_activation,
    )

forward

forward(input_tensor, offset=0)

Compute 2D positional encodings for (B, C, H, W) feature maps.

Source code in src/versatil/models/layers/positional_encoding/base.py

def forward(self, input_tensor: torch.Tensor, offset: int = 0) -> torch.Tensor:
    """Compute 2D positional encodings for (B, C, H, W) feature maps."""
    batch_size, channels, height, width = input_tensor.shape

    encodings = self._compute_encodings(
        torch.empty(height, width).to(input_tensor.device)
    )  # [embedding_dimension, height, width]
    encodings = encodings.unsqueeze(0).repeat(
        batch_size, 1, 1, 1
    )  # [batch_size, embedding_dimension, height, width]
    if self.mlp_network:
        # Reshape for MLP: [batch_size, embedding_dimension, height, width] -> [batch_size * height * width, embedding_dimension]
        encodings = encodings.permute(0, 2, 3, 1).reshape(
            -1, self.embedding_dimension
        )
        encodings = self.mlp_network(encodings)
        output_dimension = encodings.shape[-1]
        # Reshape back: [batch_size, output_dimension, height, width]
        encodings = encodings.reshape(
            batch_size, height, width, output_dimension
        ).permute(0, 3, 1, 2)
    return encodings

add_positional_encoding

add_positional_encoding(source, positional_encoding=None)

Adds positional encoding to the tensor if provided.

Parameters:

Name	Type	Description	Default
source	Tensor	Input tensor.	required
positional_encoding	Tensor | None	Positional encoding tensor to add (optional).	None

Returns:

Type	Description
Tensor	Tensor with positional encoding added if provided, otherwise the original tensor.

Source code in src/versatil/models/layers/positional_encoding/base.py

def add_positional_encoding(
    source: torch.Tensor, positional_encoding: torch.Tensor | None = None
) -> torch.Tensor:
    """Adds positional encoding to the tensor if provided.

    Args:
        source: Input tensor.
        positional_encoding: Positional encoding tensor to add (optional).

    Returns:
        Tensor with positional encoding added if provided, otherwise the original tensor.
    """
    return source if positional_encoding is None else source + positional_encoding