lgatr.layers.slim_layers.SlimMLP

class lgatr.layers.slim_layers.SlimMLP(v_channels, s_channels, nonlinearity='gelu', nonlinearity_v='sigmoid', mlp_ratio=2, num_layers=2, dropout_prob=None)

Bases: Module

Multi-layer perceptron for vector and scalar features.

Parameters:

• v_channels (int) – Number of vector channels.

• s_channels (int) – Number of scalar channels.

• nonlinearity (str) – Nonlinearity for the GLU layers (scalar gate, and vector gate when nonlinearity_v is None).

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

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

• num_layers (int) – Total number of layers (must be >= 2).

• dropout_prob (float | None) – Dropout probability.

forward(vectors, scalars)

Forward pass.

Parameters:

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

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

Return type:

tuple[Tensor, Tensor]

Returns:

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

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