File size: 1,236 Bytes
b615ad6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """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())
|