Long Short-Term Memory

"""One LSTM cell and a forward pass over a sequence - the gates, written out.

Four gates, all the same shape: forget decides what to drop from the cell
state, input decides what to consider writing, candidate is what to write,
and output decides what to expose as the hidden state.
"""

import math


def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z))


def lstm_cell(x, h_prev, c_prev, weights, hidden_size):
    """One timestep. x, h_prev, c_prev are plain lists of floats."""
    combined = list(x) + list(h_prev)

    def gate(name):
        w, b = weights[name]
        out = []
        for j in range(hidden_size):
            acc = b[j]
            for k in range(len(combined)):
                acc += w[j][k] * combined[k]
            out.append(acc)
        return out

    f = [sigmoid(v) for v in gate("forget")]     # what to erase
    i = [sigmoid(v) for v in gate("input")]      # what to admit
    g = [math.tanh(v) for v in gate("candidate")]  # what to write
    o = [sigmoid(v) for v in gate("output")]     # what to expose

    # The additive update is the whole point: c flows forward through a sum,
    # not a repeated matrix product, so its gradient does not vanish.
    c = [f[j] * c_prev[j] + i[j] * g[j] for j in range(hidden_size)]
    h = [o[j] * math.tanh(c[j]) for j in range(hidden_size)]
    return h, c


def forward(sequence, weights, hidden_size):
    h = [0.0] * hidden_size
    c = [0.0] * hidden_size
    for x in sequence:
        h, c = lstm_cell(x, h, c, weights, hidden_size)
    return h

"""LSTM forecaster - typed, validated, and expressed in a framework."""

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class ForecastConfig:
    input_size: int
    hidden_size: int = 64
    num_layers: int = 2
    horizon: int = 24
    dropout: float = 0.1


class LstmForecaster(nn.Module):
    """Encodes a window and emits all `horizon` steps at once.

    Direct multi-horizon output rather than autoregressive decoding: feeding
    predictions back compounds error across the horizon, and the head costs
    almost nothing.
    """

    def __init__(self, config: ForecastConfig) -> None:
        super().__init__()
        if config.num_layers < 1:
            raise ValueError(f"num_layers must be >= 1, got {config.num_layers}")
        if not 0.0 <= config.dropout < 1.0:
            raise ValueError(f"dropout must be in [0, 1), got {config.dropout}")

        self.config = config
        self.lstm = nn.LSTM(
            input_size=config.input_size,
            hidden_size=config.hidden_size,
            num_layers=config.num_layers,
            # Dropout between layers only. Applying it across time would sever
            # the memory the architecture exists to preserve.
            dropout=config.dropout if config.num_layers > 1 else 0.0,
            batch_first=True,
        )
        self.head = nn.Linear(config.hidden_size, config.horizon)

    def forward(self, windows: Tensor) -> Tensor:
        if windows.dim() != 3:
            raise ValueError(f"expected (batch, time, features), got {tuple(windows.shape)}")

        outputs, _ = self.lstm(windows)
        return self.head(outputs[:, -1, :])   # last step summarizes the window


def train_step(
    model: LstmForecaster,
    optimizer: torch.optim.Optimizer,
    windows: Tensor,
    targets: Tensor,
    clip_norm: float = 1.0,
) -> float:
    model.train()
    optimizer.zero_grad(set_to_none=True)
    loss = nn.functional.mse_loss(model(windows), targets)
    loss.backward()
    # Not optional hygiene: one steep batch can send the gradient norm to
    # infinity and destroy the weights irrecoverably.
    nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
    optimizer.step()
    return float(loss.item())

"""LSTM forecaster - packed batches, fused kernels, mixed precision."""

import torch
from torch import Tensor, nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence


class FastLstmForecaster(nn.Module):
    def __init__(self, input_size: int, hidden_size: int = 64, horizon: int = 24) -> None:
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers=2, batch_first=True)
        self.head = nn.Linear(hidden_size, horizon)

    def forward(self, windows: Tensor, lengths: Tensor) -> Tensor:
        # Packing skips padded timesteps entirely. On a batch bucketed by
        # length this is a large constant-factor win, and it also stops the
        # padding from polluting the final hidden state.
        packed = pack_padded_sequence(
            windows, lengths.cpu(), batch_first=True, enforce_sorted=False
        )
        packed_out, (hidden, _) = self.lstm(packed)
        pad_packed_sequence(packed_out, batch_first=True)
        return self.head(hidden[-1])


def train_epoch(model, optimizer, loader, device, clip_norm: float = 1.0) -> float:
    model.train()
    scaler = torch.amp.GradScaler(device)
    total = 0.0

    for windows, lengths, targets in loader:
        windows = windows.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)
        # bfloat16 autocast: the recurrence is memory-bandwidth bound, so
        # halving the activation traffic is close to halving the step time.
        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            loss = nn.functional.mse_loss(model(windows, lengths), targets)

        scaler.scale(loss).backward()
        # Unscale before clipping, or the threshold is applied to scaled
        # gradients and means nothing.
        scaler.unscale_(optimizer)
        nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
        scaler.step(optimizer)
        scaler.update()
        total += float(loss.item())

    return total / max(len(loader), 1)