ml.models.activations

Defines a general-purpose API for activation functions.

from ml.models.activations import get_activation, cast_activation_type

model = nn.Sequential(nn.Linear(4, 5), get_activation("relu"))

# This lets you parametrize the activation function as a string.
model = nn.Sequential(nn.Linear(4, 5), get_activation(cast_activation_type(my_activation)))

Choices for the activation functions are:

• "no_act"

• "relu"

• "relu6"

• "relu2"

• "clamp6"

• "leaky_relu"

• "elu"

• "celu"

• "selu"

• "gelu"

• "gelu_fast"

• "sigmoid"

• "log_sigmoid"

• "hard_sigomid"

• "tanh"

• "softsign"

• "softplus"

• "silu"

• "mish"

• "swish"

• "hard_swish"

• "soft_shrink"

• "hard_shrink"

• "tanh_shrink"

• "soft_sign"

• "relu_squared"

• "laplace"

ml.models.activations.cast_activation_type(s: str) → Literal['no_act', 'relu', 'relu6', 'relu2', 'clamp6', 'leaky_relu', 'elu', 'celu', 'selu', 'gelu', 'gelu_fast', 'sigmoid', 'log_sigmoid', 'hard_sigomid', 'tanh', 'softsign', 'softplus', 'silu', 'mish', 'swish', 'hard_swish', 'soft_shrink', 'hard_shrink', 'tanh_shrink', 'soft_sign', 'relu_squared', 'laplace']

class ml.models.activations.Clamp(*, value: float | None = None, value_range: tuple[float, float] | None = None, inplace: bool = False)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.activations.Clamp6(inplace: bool = False)

Bases: Clamp

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class ml.models.activations.ReLUSquared(*args, **kwargs)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.activations.FastGELU(*args, **kwargs)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.activations.QuickGELU(*args, **kwargs)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ml.models.activations.LaplaceActivation(mu: float = 0.707107, sigma: float = 0.282095)

Bases: Module

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

ml.models.activations.get_activation(act: Literal['no_act', 'relu', 'relu6', 'relu2', 'clamp6', 'leaky_relu', 'elu', 'celu', 'selu', 'gelu', 'gelu_fast', 'sigmoid', 'log_sigmoid', 'hard_sigomid', 'tanh', 'softsign', 'softplus', 'silu', 'mish', 'swish', 'hard_swish', 'soft_shrink', 'hard_shrink', 'tanh_shrink', 'soft_sign', 'relu_squared', 'laplace'], *, inplace: bool = True) → Module

Returns an activation function from a keyword string.

Parameters:

• act – The keyword for the activation function (None for identity)

• inplace – If set, use the inplace version of the activation function

Returns:

The activation function as a module

Raises:

NotImplementedError – If the activation function is invalid