spd_learn.modules.SPDTemporalEncoder#

class spd_learn.modules.SPDTemporalEncoder(sequence_length=30, embed_dim=6, num_heads=2, feedforward_dim=10, num_layers=1, dropout=0.1, *, device=None, dtype=None)[source]#

Bases: Module

Stack positional encoding and structure-preserving encoder layers.

See [Seraphim et al., 2024].

Parameters:

  • sequence_length (int, default=30) – Number of ordered input tokens.

  • embed_dim (int, default=6) – Triangular log-coordinate dimension.

  • num_heads (int, default=2) – Number of attention heads; must divide embed_dim.

  • feedforward_dim (int, default=10) – Width of the encoder feedforward layer.

  • num_layers (int, default=1) – Number of independently registered encoder layers.

  • dropout (float, default=0.1) – Dropout probability.

  • device (torch.device or str, optional) – Device for parameters and buffers.

  • dtype (torch.dtype, optional) – Data type for parameters and buffers.

forward(x: Tensor) → Tensor[source]#

Encode an ordered log-token sequence.

Parameters:

x (torch.Tensor) – Log tokens of shape (batch, tokens, embed_dim).

Returns:

Encoded tokens with the input shape.

Return type:

torch.Tensor