Reference / API reference / dorothy / dorothy.models / dorothy.models.utils

dorothy.models.utils

Shared utilities for neural network model building.

This module provides common functions used across different model architectures, including activation and normalization layer factories and the residual block of the multi-head model’s residual stream.

Functions

get_activation(activation)	Get the activation function class by name.
get_normalization(normalization)	Get the normalization layer class by name.
hidden_blocks(input_dim, hidden_layers, ...)	Build the hidden layers of a network, one block per layer.
residual_blocks(width, hidden_layers, ...)	Build a residual stream: one ResidualBlock per hidden size, then a norm.

Classes

ResidualBlock(width, hidden, normalization, ...)

A pre-norm residual block of constant width.

dorothy.models.utils.get_activation(activation)

Get the activation function class by name.

Parameters:

activation (str) – Name of activation function (“gelu”, “relu”, “silu”).

Returns:

The activation module class.

Raises:

ValueError – If activation name is not recognized.

Return type:

type[Module]

Examples

>>> act_class = get_activation("gelu")
>>> layer = act_class()

dorothy.models.utils.get_normalization(normalization)

Get the normalization layer class by name.

Parameters:

normalization (str) – Name of normalization (“batchnorm”, “layernorm”).

Returns:

The normalization module class.

Raises:

ValueError – If normalization name is not recognized.

Return type:

type[Module]

Examples

>>> norm_class = get_normalization("layernorm")
>>> layer = norm_class(256)  # 256 features

dorothy.models.utils.hidden_blocks(input_dim, hidden_layers, normalization, activation, dropout)

Build the hidden layers of a network, one block per layer.

Each block is Linear, normalization, activation and, when dropout > 0, Dropout.

Parameters:

• input_dim (int) – Size of the input.

• hidden_layers (list[int]) – Size of each hidden layer, in order.

• normalization (str) – Normalization type (“batchnorm” or “layernorm”).

• activation (str) – Activation function (“gelu”, “relu” or “silu”).

• dropout (float) – Dropout probability.

Returns:

The modules of all blocks, in order.

Raises:

ValueError – If the normalization or activation is not recognized.

Return type:

list[Module]

class dorothy.models.utils.ResidualBlock(width, hidden, normalization, activation, dropout)

Bases: Module

A pre-norm residual block of constant width.

For an input x of size width:

x <- x + W_2 act(W_1 norm(x))

W_1 maps width to hidden and W_2 maps hidden back to width, both with a bias. When dropout > 0, Dropout follows the activation.

update

The layers that compute the update added to x.

forward(x)

Add the block’s update to x.

Parameters:

x (Tensor) – Input of shape (batch, width).

Returns:

Output of shape (batch, width).

Return type:

Tensor

dorothy.models.utils.residual_blocks(width, hidden_layers, normalization, activation, dropout)

Build a residual stream: one ResidualBlock per hidden size, then a norm.

Parameters:

• width (int) – Size of the stream.

• hidden_layers (list[int]) – Hidden size of each block, in order.

• normalization (str) – Normalization type (“batchnorm” or “layernorm”).

• activation (str) – Activation function (“gelu”, “relu” or “silu”).

• dropout (float) – Dropout probability.

Returns:

The blocks and the closing normalization, in order.

Raises:

ValueError – If the normalization or activation is not recognized.

Return type:

list[Module]