"""
One transformer block: a pre-norm residual around attention, then another
around the SwiGLU MLP

    x = x + attn(norm1(x))
    x = x + mlp(norm2(x))

Pre-norm (norm before each sublayer) keeps a clean residual highway, which
lets deep stacks train stably.    
"""

import torch.nn as nn

from .norm import AbyssRMSNorm
from .attention import AbyssCausalSelfAttention
from .mlp import AbyssSwiGLU


class AbyssBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm1 = AbyssRMSNorm(config.n_embd, config.norm_eps)
        self.attn = AbyssCausalSelfAttention(config)
        self.norm2 = AbyssRMSNorm(config.n_embd, config.norm_eps)
        self.mlp = AbyssSwiGLU(config)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x