noether.modeling.modules.mlp.upactdown_mlp

Classes

UpActDownMLPConfig
UpActDownMlp	UpActDownMlp is a vanilla MLP with an up-projection followed by an GELU activation function and a

Module Contents

class noether.modeling.modules.mlp.upactdown_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

class noether.modeling.modules.mlp.upactdown_mlp.UpActDownMlp(config)

Bases: torch.nn.Module

UpActDownMlp is a vanilla MLP with an up-projection followed by an GELU activation function and a down-projection to the original input dim.

Initialize the UpActDownMlp.

Parameters:

config (UpActDownMLPConfig) – The configuration of the UpActDownMlp.

init_weights

fc1

act

fc2

reset_parameters()

Reset the parameters of the MLP with a specific initialization. Options are “torch” (i.e., default), or

“truncnormal002”.

Raises:

NotImplementedError – raised if the specified initialization is not implemented.

Return type:

None

forward(x)

Forward function of the UpActDownMlp.

Parameters:

x (torch.Tensor) – Input tensor to the MLP.

Returns:

Output tensor from the MLP.

Return type:

torch.Tensor