From Numbers to Arrays: Batches and Matrix Gradients

# Chapter 6: the Chapter 5 engine, one array at a time instead of one number at a time.
import numpy as np


def unbroadcast(grad, shape):
    """Sum grad back down to `shape`: a value copied across rows gets the sum of the rows' slopes."""
    while grad.ndim > len(shape):              # extra leading axes were added by broadcasting
        grad = grad.sum(axis=0)
    for axis, size in enumerate(shape):        # axes of size 1 were stretched
        if size == 1:
            grad = grad.sum(axis=axis, keepdims=True)
    return grad


class Tensor:
    """An array in a calculation, plus dL/d(every entry), in an array of the same shape."""

    def __init__(self, data, parents=(), op=""):
        self.data = np.asarray(data, dtype=float)
        self.grad = np.zeros_like(self.data)
        self._parents = parents
        self._op = op
        self._backward = lambda: None

    shape = property(lambda self: self.data.shape)
    __array_ufunc__ = None                     # with a NumPy array on the left (Y * P), NumPy hands over to us

    def __repr__(self):
        return f"Tensor(shape={self.shape}, op={self._op!r})"

    # --- elementwise operations: Chapter 5's slopes, applied entry by entry ----
    def __add__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(self.data + other.data, (self, other), "+")

        def _backward():
            self.grad += unbroadcast(out.grad, self.shape)
            other.grad += unbroadcast(out.grad, other.shape)
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Tensor) else Tensor(other)
        out = Tensor(self.data * other.data, (self, other), "*")

        def _backward():
            self.grad += unbroadcast(other.data * out.grad, self.shape)
            other.grad += unbroadcast(self.data * out.grad, other.shape)
        out._backward = _backward
        return out

    def __pow__(self, n):
        out = Tensor(self.data ** n, (self,), f"**{n}")

        def _backward():
            self.grad += n * self.data ** (n - 1) * out.grad
        out._backward = _backward
        return out

    def exp(self):
        out = Tensor(np.exp(self.data), (self,), "exp")

        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

    def log(self):
        out = Tensor(np.log(self.data), (self,), "log")

        def _backward():
            self.grad += out.grad / self.data
        out._backward = _backward
        return out

    def relu(self):
        out = Tensor(np.maximum(0, self.data), (self,), "relu")

        def _backward():                       # slope 1 where the input was positive, 0 elsewhere
            self.grad += (self.data > 0) * out.grad
        out._backward = _backward
        return out

    # --- the new operations: matrix product, transpose, sum ----------------------
    def __matmul__(self, other):
        out = Tensor(self.data @ other.data, (self, other), "@")

        def _backward():                       # C = A @ B:  dA = dC @ Bᵀ,  dB = Aᵀ @ dC
            self.grad += out.grad @ other.data.T
            other.grad += self.data.T @ out.grad
        out._backward = _backward
        return out

    @property
    def T(self):
        out = Tensor(self.data.T, (self,), "T")

        def _backward():
            self.grad += out.grad.T
        out._backward = _backward
        return out

    def sum(self):
        out = Tensor(self.data.sum(), (self,), "sum")

        def _backward():                       # every entry moves the total one-for-one
            self.grad += np.ones_like(self.data) * out.grad
        out._backward = _backward
        return out

    # --- built from the above, exactly as in Chapter 5 ------------------------------
    def __neg__(self): return self * -1
    def __sub__(self, other): return self + (-other)
    def __truediv__(self, other): return self * other ** -1 if isinstance(other, Tensor) else self * (1 / other)
    def __radd__(self, other): return self + other
    def __rmul__(self, other): return self * other
    def __rsub__(self, other): return Tensor(other) - self
    def __rtruediv__(self, other): return Tensor(other) * self ** -1
    def mean(self): return self.sum() / self.data.size
    def sigmoid(self): return 1 / (1 + (-self).exp())

    def backward(self):
        order, seen = [], set()

        def visit(t):
            if id(t) not in seen:
                seen.add(id(t))
                for p in t._parents:
                    visit(p)
                order.append(t)
        visit(self)
        self.grad = np.ones_like(self.data)    # dL/dL = 1
        for t in reversed(order):
            t._backward()


def count(t, seen=None):
    """How many Tensors the graph holds."""
    seen = set() if seen is None else seen
    if id(t) not in seen:
        seen.add(id(t))
        for p in t._parents:
            count(p, seen)
    return len(seen)


# === 1. All four rides in one go ==============================================
km = np.array([2, 5, 8, 12])
minutes = np.array([11, 21, 28, 42])
w, b = Tensor(3.0), Tensor(5.0)               # scalars, broadcast across the four rides
L = ((w * km + b - minutes) ** 2).mean()
L.backward()
print("rides:  loss", L.data, " dL/dw =", w.grad, " dL/db =", b.grad, f" ({count(L)} Tensors)")

# === 2. The six picnic days from Chapter 3, as one matrix =======================
X = np.array([[0.7, 0.2, 0.5],                # (sun, rain, wind), one row per day
              [0.1, 0.9, 0.3],
              [0.9, 0.0, 0.1],
              [0.4, 0.3, 0.9],
              [0.6, 0.1, 0.2],
              [0.5, 0.4, 0.4]])
Y = np.array([[1], [0], [1], [0], [1], [0]])  # 1 = the picnic went well; shape (6, 1)

W = Tensor([[2.0, -3.0, -1.0]])               # one neuron: (neurons × inputs) = (1 × 3), as in Chapter 2
b = Tensor([0.5])                             # one bias per neuron
Z = Tensor(X) @ W.T + b                       # (6 × 3)(3 × 1) + (1,) -> (6 × 1)
P = Z.sigmoid()
loss = -(Y * P.log() + (1 - Y) * (1 - P).log()).mean()
loss.backward()
print("\npicnic: mean cross-entropy", round(float(loss.data), 4))
print("        W.grad =", np.round(W.grad, 4), " b.grad =", np.round(b.grad, 4))
print("        shapes: W", W.shape, "W.grad", W.grad.shape, "| b", b.shape, "b.grad", b.grad.shape)

# Chapter 4's shortcut, now for a batch: dW = (P - Y)ᵀ X / N
print("        (P - Y)ᵀ X / N =", np.round((P.data - Y).T @ X / len(X), 4))

# === 3. Gradient check on every entry of W ======================================
def picnic_loss(Wv, bv):
    p = 1 / (1 + np.exp(-(X @ Wv.T + bv)))
    return float(-(Y * np.log(p) + (1 - Y) * np.log(1 - p)).mean())

h, numeric = 1e-5, np.zeros_like(W.data)
for j in range(3):
    up, down = W.data.copy(), W.data.copy()
    up[0, j] += h
    down[0, j] -= h
    numeric[0, j] = (picnic_loss(up, b.data) - picnic_loss(down, b.data)) / (2 * h)
rel = np.abs(W.grad - numeric) / (np.abs(W.grad) + np.abs(numeric))
print("        gradient check, worst relative error:", f"{rel.max():.1e}")

# === 4. Two layers: Chapter 2's three neurons, ReLU, then one output neuron =====
W1 = Tensor([[2.0, -3.0, -1.0],               # picnic
             [-1.0, 0.0, 3.0],                # kite flying
             [0.0, 2.5, 0.0]])                # stay in and read
b1 = Tensor([0.5, -0.5, 0.0])
W2 = Tensor([[1.0, -0.5, -1.0]])              # the output neuron weighs the three opinions
b2 = Tensor([0.0])

H = (Tensor(X) @ W1.T + b1).relu()            # (6 × 3)
P2 = (H @ W2.T + b2).sigmoid()                # (6 × 1)
loss2 = -(Y * P2.log() + (1 - Y) * (1 - P2).log()).mean()
loss2.backward()
print("\ntwo layers: loss", round(float(loss2.data), 4), f" ({count(loss2)} Tensors)")
for name, t in [("W1", W1), ("b1", b1), ("W2", W2), ("b2", b2)]:
    print(f"  {name}: data {t.shape}  grad {t.grad.shape}  grad = {np.round(t.grad, 4).tolist()}")

# === 5. The graph doesn't grow with the batch ===================================
rng = np.random.default_rng(0)
for n in [6, 1_000, 100_000]:
    Xb = rng.random((n, 3))
    Yb = (rng.random((n, 1)) < 0.5).astype(float)
    Wb, bb = Tensor([[2.0, -3.0, -1.0]]), Tensor([0.5])
    Pb = (Tensor(Xb) @ Wb.T + bb).sigmoid()
    lb = -(Yb * Pb.log() + (1 - Yb) * (1 - Pb).log()).mean()
    print(f"batch of {n:>7}: {count(lb)} Tensors in the graph")