import torch
import torch.nn as nn
import torch.nn.functional as F


class SwiGLU(nn.Module):
    """SwiGLU feed-forward block (Section 4.6), as used in Llama/PaLM.

    Standard formulation: down_proj(silu(gate_proj(x)) * up_proj(x))
    The inner dim is scaled down from the naive 4x so that SwiGLU's extra
    gate_proj matrix doesn't blow the parameter budget relative to a plain MLP
    of the same nominal "4x" size -- this matches how Llama-style models size it.
    """

    def __init__(self, hidden_dim: int, mult: float = 4.0):
        super().__init__()
        # standard correction: 4 * hidden * (2/3) keeps param count comparable
        # to a plain (non-gated) 4x MLP, rounded to a clean multiple of 8.
        inner_dim = int(hidden_dim * mult * 2 / 3)
        inner_dim = ((inner_dim + 7) // 8) * 8

        self.gate_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
        self.up_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
        self.down_proj = nn.Linear(inner_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))