Recurrent Neural Network

"""A vanilla RNN - forward through time and BPTT, transcribed.

The forward pass carries a state; the backward pass walks the same chain in
reverse, multiplying by the recurrent Jacobian once per step. That repeated
multiplication is the entire vanishing-gradient problem, and it is visible in
the loop below.
"""

import math
import random


def tanh(z):
    return math.tanh(z)


def fit(sequences, targets, hidden, lr=0.01, epochs=100, seed=0):
    """One hidden layer, tanh activation, scalar output per sequence."""
    rng = random.Random(seed)
    d = len(sequences[0][0])

    W_xh = [[rng.gauss(0.0, 0.1) for _ in range(hidden)] for _ in range(d)]
    W_hh = [[rng.gauss(0.0, 0.1) for _ in range(hidden)] for _ in range(hidden)]
    W_hy = [rng.gauss(0.0, 0.1) for _ in range(hidden)]
    b_h = [0.0] * hidden
    b_y = 0.0

    for _ in range(epochs):
        for sequence, target in zip(sequences, targets):
            T = len(sequence)

            # ---- forward, carrying the state --------------------------
            # Every state is kept: the backward pass needs all of them,
            # which is why memory grows with sequence length and why
            # truncation exists.
            states = [[0.0] * hidden]
            for t in range(T):
                current = [0.0] * hidden
                for h in range(hidden):
                    total = b_h[h]
                    for j in range(d):
                        total += sequence[t][j] * W_xh[j][h]
                    for k in range(hidden):
                        total += states[t][k] * W_hh[k][h]
                    current[h] = tanh(total)
                states.append(current)

            prediction = b_y
            for h in range(hidden):
                prediction += states[T][h] * W_hy[h]

            # ---- backward through time --------------------------------
            error = 2.0 * (prediction - target)

            grad_W_hy = [error * states[T][h] for h in range(hidden)]
            grad_b_y = error

            grad_W_xh = [[0.0] * hidden for _ in range(d)]
            grad_W_hh = [[0.0] * hidden for _ in range(hidden)]
            grad_b_h = [0.0] * hidden

            # delta at the final state, then walked backwards.
            delta = [error * W_hy[h] for h in range(hidden)]

            for t in reversed(range(T)):
                # Multiply by the activation derivative: 1 - tanh^2, which is
                # at most 1 and usually much less. That factor alone shrinks
                # the signal at every step.
                gated = [delta[h] * (1.0 - states[t + 1][h] ** 2) for h in range(hidden)]

                for h in range(hidden):
                    grad_b_h[h] += gated[h]
                    for j in range(d):
                        grad_W_xh[j][h] += gated[h] * sequence[t][j]
                    for k in range(hidden):
                        grad_W_hh[k][h] += gated[h] * states[t][k]

                # THE PRODUCT. delta is multiplied by W_hh once per step, so
                # after T steps it has been multiplied by the same matrix T
                # times - geometric decay if its spectral radius is below one,
                # geometric growth if above. No learning rate fixes either.
                delta = [
                    sum(gated[h] * W_hh[k][h] for h in range(hidden))
                    for k in range(hidden)
                ]

            # ---- update ------------------------------------------------
            for h in range(hidden):
                W_hy[h] -= lr * grad_W_hy[h]
                b_h[h] -= lr * grad_b_h[h]
                for j in range(d):
                    W_xh[j][h] -= lr * grad_W_xh[j][h]
                for k in range(hidden):
                    W_hh[k][h] -= lr * grad_W_hh[k][h]
            b_y -= lr * grad_b_y

    return W_xh, W_hh, W_hy, b_h, b_y

"""An RNN - typed, batched, orthogonally initialized, clipped, truncated."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]


@dataclass
class RecurrentLayer:
    """Weights shared across every timestep - which is what makes the
    parameter count independent of sequence length, and the gradient a
    repeated product of the same matrix."""

    W_xh: Matrix
    W_hh: Matrix
    bias: Vector

    @classmethod
    def initialized(cls, input_size: int, hidden: int, rng: np.random.Generator) -> "RecurrentLayer":
        """Orthogonal initialization for the RECURRENT matrix specifically.

        An orthogonal matrix has every singular value equal to one, so the
        repeated product neither shrinks nor grows at initialization. That
        does not solve the vanishing gradient - training moves the weights
        away from orthogonality - but it delays the onset of both failures
        long enough for learning to start, which random init frequently does
        not.
        """
        random_matrix = rng.normal(0.0, 1.0, size=(hidden, hidden))
        orthogonal, _ = np.linalg.qr(random_matrix)

        return cls(
            W_xh=rng.normal(0.0, np.sqrt(2.0 / input_size), size=(input_size, hidden)),
            W_hh=orthogonal,
            bias=np.zeros(hidden, dtype=np.float64),
        )

    def forward(self, sequence: Tensor) -> tuple[Tensor, Vector]:
        """sequence is (batch, time, features). Returns every state, because
        the backward pass needs all of them - which is what makes memory grow
        with sequence length and why truncation exists."""
        if sequence.ndim != 3:
            raise ValueError(f"expected (batch, time, features), got {sequence.shape}")
        if sequence.shape[2] != self.W_xh.shape[0]:
            raise ValueError(
                f"expected {self.W_xh.shape[0]} features, got {sequence.shape[2]}"
            )

        batch, steps, _ = sequence.shape
        hidden = self.bias.size
        states = np.zeros((batch, steps + 1, hidden), dtype=np.float64)

        # The time loop cannot be vectorized - step t reads the state step
        # t-1 produced. What CAN be batched is the batch dimension, so each
        # step is one matrix product over the whole batch rather than a loop
        # over sequences.
        for t in range(steps):
            pre_activation = (
                sequence[:, t] @ self.W_xh + states[:, t] @ self.W_hh + self.bias
            )
            states[:, t + 1] = np.tanh(pre_activation)

        return states, states[:, -1]


def clip_global_norm(gradients: list[Matrix], max_norm: float) -> float:
    """Clip by the GLOBAL norm across every parameter, not per tensor.

    Exploding gradients are the tractable half of the problem, and this is the
    standard remedy. Clipping tensors independently would rescale them by
    different factors and change the gradient DIRECTION; clipping the global
    norm rescales everything by one factor and preserves it.
    """
    total = np.sqrt(sum(float(np.sum(g * g)) for g in gradients))
    if total > max_norm:
        scale = max_norm / total
        for gradient in gradients:
            gradient *= scale
    return total


def truncated_bptt(
    layer: RecurrentLayer,
    sequence: Tensor,
    targets: Tensor,
    window: int = 64,
) -> list[tuple[Tensor, Vector]]:
    """Split the sequence into windows, carrying the state and cutting the
    gradient at each boundary.

    Bounds memory and compute - and also bounds what the model can learn to
    the window length, which is a limit chosen in advance rather than
    discovered. That trade is the whole point of the technique.
    """
    if window < 1:
        raise ValueError(f"window must be at least 1, got {window}")

    chunks: list[tuple[Tensor, Vector]] = []
    steps = sequence.shape[1]

    for start in range(0, steps, window):
        stop = min(start + window, steps)
        states, final = layer.forward(sequence[:, start:stop])
        # The carried state is DETACHED: it flows forward as a value and the
        # gradient does not flow back through it. That cut is the truncation.
        chunks.append((states, final.copy()))

    return chunks

"""An RNN - fused input projection, preallocated states, packed batches."""

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]


class FusedRnn:
    """The time loop is irreducibly sequential; everything around it is not.

    Step t reads the state step t-1 produced, so no rewrite parallelizes over
    time - that limitation is the architecture, and it is why attention
    displaced it. What CAN be lifted out of the loop is the input projection:
    x_t @ W_xh does not depend on any state, so all T of them are ONE matrix
    product computed before the loop starts, leaving only the recurrent term
    inside it.

    That halves the work in the hot loop and is the main reason a production
    RNN kernel looks nothing like the textbook formulation.
    """

    def __init__(self, input_size: int, hidden: int, batch: int, max_steps: int,
                 seed: int = 0) -> None:
        rng = np.random.default_rng(seed)

        # Contiguous float32: the recurrent product is memory-bound, and
        # halving the traffic on the state and weight arrays is worth more
        # than the extra mantissa bits on a bounded tanh output.
        self._W_xh = np.ascontiguousarray(
            rng.normal(0.0, np.sqrt(2.0 / input_size), size=(input_size, hidden)),
            dtype=np.float32,
        )
        orthogonal, _ = np.linalg.qr(rng.normal(0.0, 1.0, size=(hidden, hidden)))
        self._W_hh = np.ascontiguousarray(orthogonal, dtype=np.float32)
        self._bias = np.zeros(hidden, dtype=np.float32)

        # Every buffer allocated once for the whole run: states, the
        # projected inputs, and the per-step scratch.
        self._states = np.zeros((batch, max_steps + 1, hidden), dtype=np.float32)
        self._projected = np.zeros((batch, max_steps, hidden), dtype=np.float32)
        self._scratch = np.zeros((batch, hidden), dtype=np.float32)
        self._hidden = hidden

    def forward(self, sequence: Tensor, lengths: NDArray) -> Tensor:
        """sequence is (batch, time, features); lengths gives the true length
        of each member so padding contributes nothing."""
        batch, steps, _ = sequence.shape
        states = self._states[:batch, : steps + 1]
        states.fill(0.0)

        # ALL input projections in one GEMM, before the loop. This term does
        # not depend on the state, so computing it per step inside the loop -
        # as the textbook formulation does - repeats a matrix product T times
        # for no reason.
        projected = self._projected[:batch, :steps]
        np.matmul(sequence.reshape(-1, sequence.shape[2]), self._W_xh,
                  out=projected.reshape(-1, self._hidden))
        projected += self._bias

        for t in range(steps):
            # Only the recurrent term remains inside the loop.
            np.matmul(states[:, t], self._W_hh, out=self._scratch[:batch])
            self._scratch[:batch] += projected[:, t]
            np.tanh(self._scratch[:batch], out=states[:, t + 1])

            # Masking: a sequence that has ended keeps its final state rather
            # than continuing to consume padding. Done in place, because a
            # boolean-indexed copy here would allocate once per timestep.
            finished = lengths <= t
            if finished.any():
                states[finished, t + 1] = states[finished, t]

        return states[:, 1 : steps + 1]

    def final_states(self, lengths: NDArray) -> Matrix:
        """Gather each sequence's state at ITS OWN final step, not the last
        padded one - the classic off-by-one in variable-length batching."""
        rows = np.arange(lengths.size)
        return self._states[rows, lengths]