ResidualMLP

class sdm.nn.ResidualMLP(in_channels: int, hidden_channels: int, out_channels: int, activation: Module, bias: bool = True, device: device | str | None = None, dtype: dtype | None = None)

Bases: Module

Two-layer MLP with a learned linear residual projection.

\[\operatorname{MLP}_2(x) + \operatorname{Linear}(x)\]

Parameters:

• in_channels (int) – The number of input channels.

• hidden_channels (int) – The number of hidden channels in the MLP.

• out_channels (int) – The number of output channels.

• activation (Module) – The activation in the MLP.

• bias (bool) – Whether to learn additive biases.

• device (device | str | None) – The device.

• dtype (dtype | None) – The dtype.

forward(x: Tensor) → Tensor

The forward pass.

Parameters:

x (Tensor) – The tensor.

Return type:

Tensor