Multilayer Perceptron

"""A two-layer network - forward and backward, transcribed.

Backpropagation is the chain rule with the shared intermediates computed once
instead of once per weight. Every loop below is one factor of that chain,
written out so the dependency is visible.
"""

import math
import random


def sigmoid(z):
    return 1.0 / (1.0 + math.exp(-z)) if z >= 0 else math.exp(z) / (1.0 + math.exp(z))


def fit(X, y, hidden, lr=0.1, epochs=1000, seed=0):
    """One hidden layer, sigmoid activation, squared error."""
    rng = random.Random(seed)
    n = len(X)
    d = len(X[0])

    # Small random init, NOT zeros: with identical weights every hidden unit
    # computes the same thing and receives the same gradient forever, so the
    # layer never differentiates. Symmetry breaking is why this is random.
    W1 = [[rng.gauss(0.0, 0.1) for _ in range(hidden)] for _ in range(d)]
    b1 = [0.0] * hidden
    W2 = [rng.gauss(0.0, 0.1) for _ in range(hidden)]
    b2 = 0.0

    for _ in range(epochs):
        for i in range(n):
            # ---- forward ----------------------------------------------
            z1 = [0.0] * hidden
            a1 = [0.0] * hidden
            for h in range(hidden):
                total = b1[h]
                for j in range(d):
                    total += X[i][j] * W1[j][h]
                z1[h] = total
                a1[h] = sigmoid(total)

            z2 = b2
            for h in range(hidden):
                z2 += a1[h] * W2[h]
            prediction = z2

            # ---- backward ---------------------------------------------
            # delta at the output: dL/dz2 for squared error is 2 * error.
            delta2 = 2.0 * (prediction - y[i])

            # delta at the hidden layer: propagate through W2, then multiply
            # by the activation derivative. That elementwise product is the
            # gate - a saturated unit has derivative near zero and passes
            # nothing back, which is the vanishing-gradient problem in
            # miniature.
            delta1 = [0.0] * hidden
            for h in range(hidden):
                delta1[h] = delta2 * W2[h] * a1[h] * (1.0 - a1[h])

            # ---- update ------------------------------------------------
            for h in range(hidden):
                W2[h] -= lr * delta2 * a1[h]
            b2 -= lr * delta2

            for j in range(d):
                for h in range(hidden):
                    W1[j][h] -= lr * delta1[h] * X[i][j]
            for h in range(hidden):
                b1[h] -= lr * delta1[h]

    return W1, b1, W2, b2


def predict(W1, b1, W2, b2, x):
    hidden = len(b1)
    total = b2
    for h in range(hidden):
        z = b1[h]
        for j in range(len(x)):
            z += x[j] * W1[j][h]
        total += sigmoid(z) * W2[h]
    return total

"""An MLP - typed layers, He init, ReLU, fused softmax cross-entropy."""

from dataclasses import dataclass, field

import numpy as np
from numpy.typing import NDArray

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


@dataclass
class DenseLayer:
    """One affine map plus a nonlinearity, with its own backward pass."""

    weights: Matrix
    bias: Vector
    _inputs: Matrix | None = field(default=None, init=False, repr=False)
    _pre_activation: Matrix | None = field(default=None, init=False, repr=False)

    @classmethod
    def he_initialized(cls, fan_in: int, fan_out: int, rng: np.random.Generator) -> "DenseLayer":
        """Scale by sqrt(2 / fan_in) - He initialization for ReLU.

        Not a refinement. With the wrong scale the activation variance grows
        or shrinks by a constant factor at every layer, so a deep stack either
        saturates or collapses to zero before the first gradient arrives.
        Getting this wrong makes a network untrainable rather than slow.
        """
        scale = np.sqrt(2.0 / fan_in)
        return cls(
            weights=rng.normal(0.0, scale, size=(fan_in, fan_out)),
            bias=np.zeros(fan_out, dtype=np.float64),
        )

    def forward(self, inputs: Matrix) -> Matrix:
        if inputs.ndim != 2 or inputs.shape[1] != self.weights.shape[0]:
            raise ValueError(
                f"expected (batch, {self.weights.shape[0]}) input, got {inputs.shape}"
            )
        # Saved for the backward pass: the gradient of the weights needs the
        # incoming activation, and this is what makes activation memory - not
        # parameter count - the thing that caps batch size.
        self._inputs = inputs
        self._pre_activation = inputs @ self.weights + self.bias
        return np.maximum(self._pre_activation, 0.0)      # ReLU

    def backward(self, upstream: Matrix, learning_rate: float) -> Matrix:
        if self._inputs is None or self._pre_activation is None:
            raise RuntimeError("backward called before forward")

        # The ReLU derivative is a gate: pass the signal where the unit was
        # active, block it where it was not. Unlike sigmoid it does not shrink
        # the signal, which is why deep stacks became trainable.
        delta = upstream * (self._pre_activation > 0.0)

        grad_weights = self._inputs.T @ delta
        grad_bias = delta.sum(axis=0)
        downstream = delta @ self.weights.T

        self.weights -= learning_rate * grad_weights
        self.bias -= learning_rate * grad_bias
        return downstream


def softmax_cross_entropy(logits: Matrix, labels: NDArray) -> tuple[float, Matrix]:
    """Loss and its gradient, computed together and stably.

    Two things matter here. The max is subtracted before exponentiating,
    because exp of a large logit overflows and the naive version silently
    produces NaN. And the gradient of the fused pair is simply
    (softmax - one_hot) - the Jacobian of the softmax and the derivative of
    the log cancel exactly, so computing them separately does more work and
    is less stable.
    """
    shifted = logits - logits.max(axis=1, keepdims=True)
    exponentials = np.exp(shifted)
    partition = exponentials.sum(axis=1, keepdims=True)

    log_probabilities = shifted - np.log(partition)
    rows = np.arange(labels.size)
    loss = float(-log_probabilities[rows, labels].mean())

    gradient = exponentials / partition
    gradient[rows, labels] -= 1.0
    return loss, gradient / labels.size


def fit(
    X: Matrix,
    y: NDArray,
    hidden_sizes: list[int],
    n_classes: int,
    learning_rate: float = 0.01,
    epochs: int = 50,
    batch_size: int = 64,
    seed: int = 0,
) -> list[DenseLayer]:
    """Train by mini-batch SGD. Raises ValueError on malformed input."""
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if X.shape[0] != y.size:
        raise ValueError(f"X has {X.shape[0]} rows but y has {y.size}")
    if batch_size < 1:
        raise ValueError(f"batch_size must be at least 1, got {batch_size}")

    rng = np.random.default_rng(seed)
    sizes = [X.shape[1], *hidden_sizes, n_classes]
    layers = [
        DenseLayer.he_initialized(sizes[i], sizes[i + 1], rng)
        for i in range(len(sizes) - 1)
    ]

    for _ in range(epochs):
        order = rng.permutation(X.shape[0])

        for start in range(0, X.shape[0], batch_size):
            batch = order[start : start + batch_size]
            activations = X[batch]

            for layer in layers:
                activations = layer.forward(activations)

            _, gradient = softmax_cross_entropy(activations, y[batch])

            for layer in reversed(layers):
                gradient = layer.backward(gradient, learning_rate)

    return layers

"""An MLP - preallocated activations, in-place ReLU, Adam in place."""

import numpy as np
from numpy.typing import NDArray

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


class FusedMlp:
    """Every buffer allocated once for the whole run.

    Training allocates the same shapes thousands of times: activations,
    gradients, and one moment estimate per parameter. At a fixed batch size
    every one of those shapes is known up front, so they are allocated once
    and rewritten - which removes the allocator from the inner loop entirely
    and keeps the working set resident.

    Activation memory, not parameter count, is what caps batch size in a real
    network. Preallocating makes that ceiling explicit rather than discovering
    it as an out-of-memory error mid-epoch.
    """

    def __init__(self, sizes: list[int], batch_size: int, seed: int = 0) -> None:
        rng = np.random.default_rng(seed)
        self._sizes = sizes
        self._batch = batch_size

        # C-contiguous float32: half the memory traffic of float64, and every
        # GEMM reads the arrays directly rather than through a copy. Training
        # is memory-bound far more often than precision-bound.
        self._weights = [
            np.ascontiguousarray(
                rng.normal(0.0, np.sqrt(2.0 / sizes[i]), size=(sizes[i], sizes[i + 1])),
                dtype=np.float32,
            )
            for i in range(len(sizes) - 1)
        ]
        self._biases = [np.zeros(sizes[i + 1], dtype=np.float32) for i in range(len(sizes) - 1)]

        # Adam moments, allocated once alongside the parameters.
        self._m = [np.zeros_like(w) for w in self._weights]
        self._v = [np.zeros_like(w) for w in self._weights]
        self._step = 0

        # Forward activations and backward deltas, one buffer per layer.
        self._activations = [
            np.empty((batch_size, size), dtype=np.float32) for size in sizes
        ]
        self._deltas = [
            np.empty((batch_size, size), dtype=np.float32) for size in sizes[1:]
        ]

    def forward(self, batch: Matrix) -> Matrix:
        np.copyto(self._activations[0], batch)

        for index, (weights, bias) in enumerate(zip(self._weights, self._biases)):
            # matmul straight into the preallocated buffer - no temporary per
            # layer per step, which across an epoch is thousands of arrays.
            np.matmul(self._activations[index], weights, out=self._activations[index + 1])
            self._activations[index + 1] += bias

            if index < len(self._weights) - 1:
                # ReLU in place: the output buffer IS the input buffer.
                np.maximum(self._activations[index + 1], 0.0, out=self._activations[index + 1])

        return self._activations[-1]

    def backward(self, gradient: Matrix, learning_rate: float,
                 beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8) -> None:
        self._step += 1
        np.copyto(self._deltas[-1], gradient)

        for index in reversed(range(len(self._weights))):
            delta = self._deltas[index]

            if index < len(self._weights) - 1:
                # ReLU gate applied in place against the saved pre-activation.
                np.multiply(delta, self._activations[index + 1] > 0.0, out=delta)

            grad_weights = self._activations[index].T @ delta
            grad_bias = delta.sum(axis=0)

            if index > 0:
                np.matmul(delta, self._weights[index].T, out=self._deltas[index - 1])

            # Adam, entirely in place: three full-size temporaries per
            # parameter tensor per step would otherwise be allocated, and with
            # a large model that is the dominant allocation in training.
            m, v = self._m[index], self._v[index]
            m *= beta1
            m += (1.0 - beta1) * grad_weights
            v *= beta2
            v += (1.0 - beta2) * (grad_weights * grad_weights)

            bias1 = 1.0 - beta1**self._step
            bias2 = 1.0 - beta2**self._step
            self._weights[index] -= learning_rate * (m / bias1) / (np.sqrt(v / bias2) + eps)
            self._biases[index] -= learning_rate * grad_bias