Activation Function Factory

Activation function factory.

Selects the appropriate activation module from the model configuration, mirroring the normalization factory pattern in src/model/norm/factory.py.

src.model.activation_function.factory.get_activation(config: Any, dim: int | None = None) → torch.nn.Module[source]

Factory function that returns an activation module based on config.

Dispatches on config.ffn_activation. Elementwise activations return a ready-to-call nn.Module. GLU variants ("swiglu", "geglu", "reglu") are not built here (they require linear projections); callers should check GLU_VARIANTS and build a GatedFFN instead.

Parameters:

  • config – Model configuration object exposing ffn_activation (string) and an optional ffn_activation_config mapping with keys such as prelu_init, elu_alpha, swish_beta, leaky_relu_slope, and the RAF keys raf_degrees, raf_version, raf_approx_func, raf_trainable, raf_input_scaling.

  • dim – Feature dimension that the activation operates on. When None, defaults to config.hidden_size (or config.dim). For FFN activations this must be the intermediate dimension (ffn_hidden_size), since the activation runs after the up-projection. Default: None.

Returns:

An nn.Module implementing the requested activation.

Raises:

ValueError – If ffn_activation is a GLU variant (callers must build the gated FFN separately) or is otherwise not recognized.