lgatr.layers.attention.self_attention.SelfAttention

class lgatr.layers.attention.self_attention.SelfAttention(config, primitives)[source]

Bases: Module

L-GATr self-attention.

Constructs queries, keys, and values, computes geometric attention, and projects linearly to outputs.

Parameters:
forward(multivectors, scalars=None, additional_qk_features_mv=None, additional_qk_features_s=None, **attn_kwargs)[source]

Compute self-attention.

Queries, keys, and values are computed per head, geometric attention is applied over items, the heads are concatenated, and a final linear map produces the outputs.

Parameters:

  • multivectors (Tensor) – Input multivectors of shape (..., items, mv_channels, 16).

  • scalars (Tensor | None) – Optional input scalars of shape (..., items, s_channels). If None, the scalar stream is bypassed and outputs_s may be None (or a tensor lifted by out_linear if out_s_channels is configured).

  • additional_qk_features_mv (Tensor | None) – Additional multivector Q/K features of shape (..., items, add_qk_mv_channels, 16).

  • additional_qk_features_s (Tensor | None) – Additional scalar Q/K features of shape (..., items, add_qk_s_channels).

  • **attn_kwargs – Optional keyword arguments forwarded to attention.

Return type:

tuple[Tensor, Tensor | None]

Returns:

  • outputs_mv – Output multivectors of shape (..., items, out_mv_channels, 16).

  • outputs_s – Output scalars of shape (..., items, out_s_channels), or None.