lgatr.layers.conditional_lgatr_block.ConditionalLGATrBlock

class lgatr.layers.conditional_lgatr_block.ConditionalLGATrBlock(mv_channels, s_channels, mv_channels_cond, s_channels_cond, attention, crossattention, mlp, primitives, dropout_prob=None, norm_elementwise_affine=True)[source]

Bases: Module

L-GATr decoder block.

Inputs are first processed by LayerNorm, multi-head geometric self-attention, and a residual connection. Then the conditions are mixed in via cross-attention with the same overhead as self-attention. Finally the data goes through 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.

  • mv_channels_cond (int) – Number of condition multivector channels.

  • s_channels_cond (int) – Number of condition scalar channels. Use 0 for no scalar condition stream.

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

  • crossattention (CrossAttentionConfig | Mapping) – Cross-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 block EquiLayerNorm instances learn an affine gain.

forward(multivectors, multivectors_cond, scalars=None, scalars_cond=None, attn_kwargs=None, crossattn_kwargs=None)[source]

Forward pass of the decoder block.

Parameters:

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

  • multivectors_cond (Tensor) – Condition multivectors of shape (..., items_cond, mv_channels_cond, 16).

  • scalars (Tensor | None) – Optional input scalars of shape (..., items, s_channels). If None, the scalar stream is bypassed and outputs_s is None.

  • scalars_cond (Tensor | None) – Optional condition scalars of shape (..., items_cond, s_channels_cond).

  • attn_kwargs (dict | None) – Optional keyword arguments forwarded to self-attention (e.g. attention masks).

  • crossattn_kwargs (dict | None) – Optional keyword arguments forwarded to cross-attention (e.g. attention masks).

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.