linmult.core.pe

Sinusoidal positional encoding with optional dropout.

Classes

PositionalEncoding

Sinusoidal positional encoding for sequence inputs.

Module Contents

class linmult.core.pe.PositionalEncoding(dropout: float = 0.1)

Bases: torch.nn.Module

Sinusoidal positional encoding for sequence inputs.

Adds fixed sinusoidal position encodings to the input tensor, following Vaswani et al. (2017). The encoding matrix is computed lazily and cached; it is only recomputed when the sequence is longer or the feature dimension changes.

Parameters:

dropout (float) – Dropout probability applied after adding the encoding. Defaults to 0.1.

Initialize PositionalEncoding.

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

Add sinusoidal positional encoding to the input.

The encoding matrix is rebuilt only when the cached tensor is shorter than the current sequence or the feature dimensionality has changed. For odd feature dimensions, the cosine slot count is floor(F/2) while the sine slot count is ceil(F/2); the division term is sliced accordingly so no index is out of range.

During tracing (torch.jit.trace/torch.onnx.export), the cache is always rebuilt instead of reused. self.pe is a single buffer mutated across repeated calls (e.g. once for the query, once for the key/value in cross-modal attention); a tracer only observes whichever branch fires on the shapes it is given, so if two calls happen to see the same time_dim at trace time, the graph bakes in one shared slice for both — which then fails at inference for genuinely different query/key lengths, even though eager mode handles that case correctly. Always rebuilding under tracing removes the shared mutable state that causes this.

Parameters:

x (torch.Tensor) – Input tensor of shape (B, T, F).

Returns:

Encoded tensor of shape (B, T, F) with dropout applied.

Return type:

torch.Tensor