spd_learn.modules.SPDInterEpochEncoder#

class spd_learn.modules.SPDInterEpochEncoder(context_epochs=21, feature_tokens=10, embed_dim=351, num_heads=3, feedforward_dim=1035, num_layers=4, dropout=0.1, *, device=None, dtype=None)[source]#

Bases: Module

Encode log tokens jointly across a chronological context.

See [Seraphim et al., 2024].

Parameters:

  • context_epochs (int, default=21) – Odd number of chronological epochs in a context.

  • feature_tokens (int, default=10) – Number of pooled tokens per epoch; must divide the input count.

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

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

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

  • num_layers (int, default=4) – 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.

Notes

Positions follow the flattened epoch-major order. Future epochs may affect every output.

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

Encode a complete epoch-major context.

Parameters:

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

Returns:

Encoded tokens with the input shape.

Return type:

torch.Tensor