spd_learn.modules.SPDTransformerEncoderLayer#

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

Bases: Module

Encode log tokens with SP-MHA and a residual feedforward block.

See [Seraphim et al., 2024].

Parameters:

  • 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.

  • 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, *, attn_mask=None, key_padding_mask=None)[source]#

Encode an ordered log-token sequence.

Parameters:

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

  • attn_mask (None) – Reserved argument; masks are not supported.

  • key_padding_mask (None) – Reserved argument; padding masks are not supported.

Returns:

Encoded tokens with the input shape.

Return type:

torch.Tensor