lgatr.layers.mlp.nonlinearities.ScalarGatedNonlinearity

class lgatr.layers.mlp.nonlinearities.ScalarGatedNonlinearity(nonlinearity='gelu')

Bases: Module

Gated nonlinearity on multivectors.

Given multivector input x, computes f(x_0) * x, where f is one of ReLU, sigmoid, tanh, GeLU, or SiLU. Auxiliary scalar inputs are processed with the same f directly (without gating).

Parameters:

nonlinearity (str) – Non-linearity type. One of "relu", "sigmoid", "tanh", "gelu", "silu".

forward(multivectors, scalars=None)

Apply the gated nonlinearity.

Parameters:

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

• scalars (Tensor | None) – Optional input scalars of shape (..., s_channels). If None, outputs_s is None.

Return type:

tuple[Tensor, Tensor | None]

Returns:

• outputs_mv – Output multivectors of shape (..., 16).

• outputs_s – Output scalars of shape (..., s_channels), or None if scalars is None.