noether.core.schemas.modules.mlp

Classes

MLPConfig
UpActDownMLPConfig

Module Contents

class noether.core.schemas.modules.mlp.MLPConfig(/, **data)

Bases: pydantic.BaseModel

Parameters:

data (Any)

input_dim: int = None

Input dimension of the MLP.

output_dim: int = None

Output dimension of the MLP.

hidden_dim: int = None

Hidden dimension for each layer.

num_layers: int = 0

Number of hidden layers in the MLP. If 0, the MLP is a two linear layer MLP from input_dim, hidden_dim, activation to output_dim.

activation: str = 'GELU'

Activation function to use between layers.

init_weights: noether.core.types.InitWeightsMode = 'truncnormal002'

Weight initialization method.

class noether.core.schemas.modules.mlp.UpActDownMLPConfig(/, **data)

Bases: pydantic.BaseModel

Parameters:

data (Any)

input_dim: int = None

Input dimension of the MLP.

hidden_dim: int = None

Hidden dimension of the MLP.

bias: bool = None

Whether to use bias in the MLP.

init_weights: noether.core.types.InitWeightsMode = None

Initialization method of the weights of the MLP. Options are “torch” (i.e., similar to the module) or ‘truncnormal002’. Defaults to ‘truncnormal002’.

check_dims()

Validator to check that hidden_dim is greater than input_dim.

Raises:

ValueError – raised if hidden_dim is not greater than input_dim.

Return type:

UpActDownMLPConfig