emmi.modules.mlp.upactdown_mlp

Classes

UpActDownMlp

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

Module Contents

class emmi.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:
  • input_dim

  • hidden_dim

  • init_weights

  • config (emmi.schemas.modules.mlp.UpActDownMLPConfig)

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