lgatr.layers.slim_layers.SlimBlock

class lgatr.layers.slim_layers.SlimBlock(v_channels, s_channels, num_heads, nonlinearity='gelu', nonlinearity_v='sigmoid', mlp_ratio=2, attn_ratio=1, num_layers_mlp=2, dropout_prob=None, norm_elementwise_affine=True)

Bases: Module

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

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

Parameters:

• v_channels (int) – Number of vector channels.

• s_channels (int) – Number of 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 the MLP’s GLUs.

• 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 RMS norms learn a per-channel gain.

forward(vectors, scalars, **attn_kwargs)

Forward pass.

Parameters:

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

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

• **attn_kwargs – Optional keyword arguments forwarded to 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).