lgatr.layers.slim_layers.SlimGLU

class lgatr.layers.slim_layers.SlimGLU(in_v_channels, out_v_channels, in_s_channels, out_s_channels, nonlinearity='gelu', nonlinearity_v='sigmoid')

Bases: Module

Gated linear unit (GLU) for vector and scalar features.

Scalar gates are computed from scalar features; vector gates are computed from inner products of (transformed) vector features.

Parameters:

• in_v_channels (int) – Number of input vector channels.

• out_v_channels (int) – Number of output vector channels.

• in_s_channels (int) – Number of input scalar channels.

• out_s_channels (int) – Number of output scalar channels.

• nonlinearity (str) – Nonlinearity for the scalar gate (and for the vector gate when nonlinearity_v is None). One of "relu", "sigmoid", "tanh", "gelu", "silu".

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

forward(vectors, scalars)

Apply the GLU.

Parameters:

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

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

Return type:

tuple[Tensor, Tensor]

Returns:

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

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