SwiGLU

class sdm.nn.SwiGLU(channels: int, hidden_channels: int, bias: bool = True, device: device | str | None = None, dtype: dtype | None = None)

Bases: Module

\(\mathrm{FFN}_{\mathrm{SwiGLU}}\) block from the “GLU Variants Improve Transformer” paper.

\[W_{\mathrm{down}}( \operatorname{Swish}_1(W_{\mathrm{gate}} x + b_{\mathrm{gate}}) \otimes (W_{\mathrm{up}} x + b_{\mathrm{up}}) ) + b_{\mathrm{down}}\]

Parameters:

• channels (int) – The number of input and output channels.

• hidden_channels (int) – The hidden channels of the gate and up projections.

• bias (bool) – If set to False, the module will not learn an additive bias.

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

• dtype (dtype | None) – The dtype.

forward(tensor: Tensor) → Tensor

The forward pass.

Parameters:

tensor (Tensor) – The tensor.

Return type:

Tensor