lgatr.layers.lgatr_block.LGATrBlock

class lgatr.layers.lgatr_block.LGATrBlock(mv_channels, s_channels, attention, mlp, primitives, dropout_prob=None, norm_elementwise_affine=True)

Bases: Module

L-GATr encoder block.

Inputs are first processed by LayerNorm, multi-head geometric self-attention, and a residual connection. Then the data is processed by another LayerNorm, an item-wise geometric MLP, and another residual connection.

Parameters:

• mv_channels (int) – Number of input and output multivector channels.

• s_channels (int) – Number of input and output scalar channels.

• attention (SelfAttentionConfig | Mapping) – Self-attention configuration.

• mlp (MLPConfig | Mapping) – MLP configuration.

• primitives (PrimitivesConfig | Mapping) – LGATr primitives configuration.

• dropout_prob (float | None) – Dropout probability.

• norm_elementwise_affine (bool) – Whether the EquiLayerNorm instances learn an affine gain.

forward(multivectors, scalars=None, additional_qk_features_mv=None, additional_qk_features_s=None, **attn_kwargs)

Forward pass of the encoder block.

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 is None.

• 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, mv_channels, 16).

• outputs_s – Output scalars of shape (..., items, s_channels), or None if scalars is None.