"""Llama 风格组件的最小可执行实现。"""
import torch
from torch import nn


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__(); self.weight = nn.Parameter(torch.ones(dim)); self.eps = eps
    def forward(self, x):
        x32 = x.float()
        return (x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)).to(x.dtype) * self.weight


def rotate_half(x):
    a, b = x[..., 0::2], x[..., 1::2]
    return torch.stack((-b, a), -1).flatten(-2)


def apply_rope(q, k, positions, base=10000.0):
    d = q.size(-1); assert d % 2 == 0
    inv = base ** (-torch.arange(0, d, 2, device=q.device).float() / d)
    angle = torch.outer(positions.float(), inv)
    cos = torch.repeat_interleave(angle.cos(), 2, -1)[None, None]
    sin = torch.repeat_interleave(angle.sin(), 2, -1)[None, None]
    return q*cos + rotate_half(q)*sin, k*cos + rotate_half(k)*sin


def repeat_kv(x, repeat):
    b, h, t, d = x.shape
    return x[:, :, None].expand(b, h, repeat, t, d).reshape(b, h*repeat, t, d)


class SwiGLU(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__(); self.g=nn.Linear(dim,hidden,bias=False); self.u=nn.Linear(dim,hidden,bias=False); self.d=nn.Linear(hidden,dim,bias=False)
    def forward(self, x): return self.d(torch.nn.functional.silu(self.g(x))*self.u(x))


if __name__ == "__main__":
    q = torch.randn(2, 4, 8, 16); k = torch.randn_like(q)
    qr, kr = apply_rope(q, k, torch.arange(8))
    assert torch.allclose(q.norm(dim=-1), qr.norm(dim=-1), atol=1e-5)
    assert repeat_kv(torch.randn(2, 2, 8, 16), 2).shape == (2, 4, 8, 16)
    assert SwiGLU(32, 64)(torch.randn(2, 5, 32)).shape == (2, 5, 32)
    print("llama component tests passed")