Reference / API reference / dorothy / dorothy.models / dorothy.models.output

dorothy.models.output

The layers that map an output head’s input to the mean and log-scatter.

Functions

output_modules(input_dim, n_parameters, ...)

Build the modules of an output head, from its input to [mu, ln s].

Classes

AdditiveOutput(input_dim, n_parameters)	The last layer of a model without label-aware uncertainty.
LabelAwareOutput(input_dim, n_parameters, ...)	Predict the mean, and ln s also from label errors, with parallel decoders.

class dorothy.models.output.LabelAwareOutput(input_dim, n_parameters, hidden_layers, normalization, activation, dropout, residual=False)

Bases: Module

Predict the mean, and ln s also from label errors, with parallel decoders.

Two decoders read the head input x. For the reported label errors sigma (normalized units), per parameter p:

mu = D_mu(x)
u_p = clamp(ln sigma_p, LN_S_MIN, LN_S_MAX) if sigma_p > 0, else m_p
ln_s = D_s([x, u])

where m is the learned missing vector. Each decoder is the hidden blocks of hidden_layers (Linear, normalization, activation, dropout), then a linear map; u enters the first layer of D_s. At least one hidden layer makes ln s a nonlinear function of u. No layer that feeds mu reads u.

With residual, x is a residual stream of width input_dim. Each decoder is one ResidualBlock per hidden size, a normalization and a linear map, and u is added to a copy of x instead of concatenated:

ln_s = D_s(x + W_u u)

mean

The label decoder D_mu.

missing_log_error

The missing vector m, shape (n_parameters,), in the units of u.

scatter

The error decoder D_s.

error_projection

The map W_u from u to the stream, without bias; None without residual.

forward(x, label_errors=None)

Predict the mean and log-scatter of every parameter.

Parameters:

• x (Tensor) – Head input, shape (batch, input_dim).

• label_errors (Tensor | None) – Reported label errors in normalized units, shape (batch, n_parameters). An entry > 0 is a reported error; any other value (0, negative, NaN) is missing. None means every entry is missing.

Returns:

Tensor of shape (batch, 2, n_parameters) = [mu, ln_s], in normalized units.

Raises:

ValueError – If label_errors does not have shape (batch, n_parameters).

Return type:

Tensor

class dorothy.models.output.AdditiveOutput(input_dim, n_parameters)

Bases: Linear

The last layer of a model without label-aware uncertainty.

One linear map from h to [mu, ln s]. Such models add the label errors to s outside the network, as sqrt(sigma_y^2 + s^2).

forward(h, label_errors=None)

Predict the mean and log-scatter of every parameter.

Parameters:

• h (Tensor) – Hidden vector, shape (batch, input_dim).

• label_errors (Tensor | None) – Must be None, since the layer reads no label errors.

Returns:

Tensor of shape (batch, 2, n_parameters) = [mu, ln_s], in normalized units.

Raises:

ValueError – If label_errors is given.

Return type:

Tensor

dorothy.models.output.output_modules(input_dim, n_parameters, hidden_layers, normalization, activation, dropout, label_aware, residual=False)

Build the modules of an output head, from its input to [mu, ln s].

Parameters:

• input_dim (int) – Size of the head input.

• n_parameters (int) – Number of predicted parameters.

• hidden_layers (list[int]) – Hidden layer sizes between the input and the means.

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

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

• dropout (float) – Dropout probability of the hidden blocks.

• label_aware (bool) – Build a LabelAwareOutput; False builds the hidden blocks and an AdditiveOutput.

• residual (bool) – Read a residual stream: the hidden layers are residual blocks of width input_dim (see LabelAwareOutput).

Returns:

The modules in order. The last one maps (features, label_errors) to (batch, 2, n_parameters); the ones before it map the input to the features.

Raises:

ValueError – If label_aware and hidden_layers is empty, or the normalization or activation is not recognized.

Return type:

list[Module]