SwiGLU

def _swish(x: Tensor) -> Tensor:
    """Swish activation σ(x) = x · sigmoid(x) (Ramachandran et al., 2018).

    Pre:
        - x is any float tensor.
    Post:
        - same shape as x with smooth non-linearity applied.
    Complexity:
        - O(1) per element.
    """
    return x * torch.sigmoid(x)


class SwiGLU(nn.Module):
    """SwiGLU feedforward block (Shazeer, 2020; Chowdhery et al., 2022).

    Math:
        SwiGLU(x) = W_o[(W_g x) ⊙ swish(W_v x)]
        where W_g splits into gate/value projections and ⊙ is element-wise multiply.

    Pre:
        - x.shape[-1] == in_features.
    Post:
        - returns a tensor with shape (..., out_features).
    Complexity:
        - O(in_features * hidden_features) per token due to the linear layers.
    Invariants:
        - gate/value split always halves the projected dimension (validated via chunk).
    """

    def __init__(
        self,
        in_features: int,
        hidden_features: int,
        out_features: int | None = None,
        bias: bool = True,
    ) -> None:
        super().__init__()
        if in_features <= 0 or hidden_features <= 0:
            raise ValueError("in_features and hidden_features must be positive.")
        self.in_features = in_features
        self.hidden_features = hidden_features
        self.out_features = out_features or in_features

        # Combined gate and value projection (efficiency)
        self.gate = nn.Linear(in_features, hidden_features * 2, bias=bias)
        self.proj = nn.Linear(hidden_features, self.out_features, bias=bias)

    def forward(self, x: Tensor) -> Tensor:
        if x.shape[-1] != self.in_features:
            raise ValueError(
                f"Input last dimension mismatch: expected {self.in_features}, got {x.shape[-1]}"
            )
        # Split gate and value from combined projection
        gate_out, value = self.gate(x).chunk(2, dim=-1)
        
        # Apply Swish to gate
        activated = _swish(gate_out)
        
        # Element-wise gating
        gated = activated * value
        
        # Project to output
        return self.proj(gated)