nessai.flows.nets

Neural networks for use in flows.

Module Contents

Classes

MLP

A standard multi-layer perceptron.

class nessai.flows.nets.MLP(in_shape, out_shape, hidden_sizes, activation=F.relu, activate_output=False, dropout_probability=0.0)

Bases: torch.nn.Module

A standard multi-layer perceptron.

Based on the implementation in nflows and modified to include dropout and conditional inputs.

Parameters:
in_shapetuple

Input shape.

out_shapetuple

Output shape.

hidden_sizesList[int]

Number of neurons in the hidden layers.

activationCallable

Activation function

activate_outputUnion[bool, Callable]

Whether to activate the output layer. If a bool is specified the same activation function is used. If a callable inputs is specified, it will be used for the activation.

dropout_probabilityfloat

Amount of dropout to apply after the hidden layers.

forward(inputs, context=None)

Forward method that allows for kwargs such as context.

Parameters:
inputstorch.tensor

Inputs to the MLP

contextNone

Conditional inputs, must be None. Only implemented to the function is compatible with other methods.

Raises:
ValueError

If the context is not None.

ValueError

If the input shape is incorrect.