layer_2.py

"""Layer 2: one recurrent block (pre-LN + residual, LSTM or GRU).

The 12-layer stack instantiates this 9 times: layers 2..10.
Causal by construction - no mask needed; memory is O(T), not O(T^2).
Hidden state resets each forward (windowed training); stateful sampling
can pass an initial state later.

Independence: imports only torch. Knows nothing about other layers.
"""
from dataclasses import dataclass

import torch.nn as nn


@dataclass
class Layer2Config:
    d_model: int = 512
    rnn_type: str = "gru"    # "lstm" | "gru"
    dropout: float = 0.0


class Layer2(nn.Module):
    def __init__(self, cfg: Layer2Config):
        super().__init__()
        self.cfg = cfg
        self.ln = nn.LayerNorm(cfg.d_model)
        if cfg.rnn_type == "lstm":
            self.rnn = nn.LSTM(cfg.d_model, cfg.d_model, batch_first=True)
        elif cfg.rnn_type == "gru":
            self.rnn = nn.GRU(cfg.d_model, cfg.d_model, batch_first=True)
        else:
            raise ValueError(f"unknown rnn_type: {cfg.rnn_type}")

    def forward(self, x, state=None):
        out, state = self.rnn(self.ln(x), state)
        return x + out, state

    @property
    def n_params(self):
        return sum(p.numel() for p in self.parameters())