lgatr.layers.slim_layers.ConditionalSlimBlock

class lgatr.layers.slim_layers.ConditionalSlimBlock(v_channels, v_channels_cond, s_channels, s_channels_cond, num_heads, nonlinearity='gelu', nonlinearity_v='sigmoid', mlp_ratio=2, attn_ratio=1, num_layers_mlp=2, dropout_prob=None, norm_elementwise_affine=True)[source]

Bases: Module

A single block of the conditional L-GATr-slim network.

Pre-norm + self-attention + residual, then pre-norm + cross-attention + residual, then pre-norm + MLP + residual.

Parameters:

  • v_channels (int) – Number of vector channels.

  • v_channels_cond (int) – Number of condition vector channels.

  • s_channels (int) – Number of scalar channels.

  • s_channels_cond (int) – Number of condition scalar channels.

  • num_heads (int) – Number of attention heads.

  • nonlinearity (str) – Nonlinearity for the MLP layers.

  • nonlinearity_v (str | None) – Optional override for the vector-path gate nonlinearity in every GLU. None falls back to nonlinearity.

  • mlp_ratio (int) – Expansion ratio for MLP hidden channels.

  • attn_ratio (int) – Expansion ratio for attention hidden channels.

  • num_layers_mlp (int) – Number of layers in the MLP (must be >= 2).

  • dropout_prob (float | None) – Dropout probability.

  • norm_elementwise_affine (bool) – Whether the SlimRMSNorm instances learn per-channel gains.

forward(vectors, vectors_cond, scalars, scalars_cond, attn_kwargs=None, crossattn_kwargs=None)[source]

Forward pass.

Parameters:

  • vectors (Tensor) – Lorentz vectors of shape (..., items, 4, v_channels).

  • vectors_cond (Tensor) – Condition Lorentz vectors of shape (..., items_cond, 4, v_channels_cond).

  • scalars (Tensor) – Scalar features of shape (..., items, s_channels).

  • scalars_cond (Tensor) – Condition scalar features of shape (..., items_cond, s_channels_cond).

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

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

Return type:

tuple[Tensor, Tensor]

Returns:

  • outputs_v – Lorentz vectors of shape (..., items, 4, v_channels).

  • outputs_s – Scalar features of shape (..., items, s_channels).