spd_learn.modules.SPDIntraEpochEncoder#

class spd_learn.modules.SPDIntraEpochEncoder(tokens_per_epoch=210, feature_tokens=10, embed_dim=351, num_heads=9, feedforward_dim=1035, num_layers=5, dropout=0.1, *, device=None, dtype=None)[source]#

Bases: Module

Encode and pool contiguous tokens within each epoch.

See [Seraphim et al., 2024].

Parameters:

  • tokens_per_epoch (int, default=210) – Number of band/window tokens per epoch.

  • 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=9) – Number of attention heads; must divide embed_dim.

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

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

feature_tokens must divide tokens_per_epoch. Epochs are encoded independently before equal-bin averaging.

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

Encode each epoch and average contiguous token groups.

Parameters:

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

Returns:

Features of shape (batch, epochs, feature_tokens, embed_dim).

Return type:

torch.Tensor