Let the Computer Do the Calculus: A Tiny Autograd Engine

# Chapter 5: a tiny autograd engine. Every number remembers how it was made,
# so the chain rule can be run backwards automatically.
# Same design as Andrej Karpathy's micrograd (MIT licence), written fresh for this course.
import math


class Value:
    """One number in a calculation, plus the slope of the final loss with respect to it."""

    def __init__(self, data, parents=(), op=""):
        self.data = float(data)
        self.grad = 0.0                # dL/d(this value), filled in by backward()
        self._parents = parents        # the Values this one was computed from
        self._op = op                  # for printing only
        self._backward = lambda: None  # pushes self.grad down to the parents

    def __repr__(self):
        return f"Value(data={self.data:.4g}, grad={self.grad:.4g})"

    # --- each operation computes its result AND records its local slopes -------
    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")

        def _backward():               # d(a + b)/da = 1, d(a + b)/db = 1
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

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

        def _backward():               # d(a·b)/da = b, d(a·b)/db = a
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

    def __pow__(self, n):              # n is a plain number, not a Value
        out = Value(self.data ** n, (self,), f"**{n}")

        def _backward():               # d(aⁿ)/da = n·aⁿ⁻¹
            self.grad += n * self.data ** (n - 1) * out.grad
        out._backward = _backward
        return out

    def exp(self):
        out = Value(math.exp(self.data), (self,), "exp")

        def _backward():               # d(eᵃ)/da = eᵃ, which is out.data
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

    def log(self):
        out = Value(math.log(self.data), (self,), "log")

        def _backward():               # d(ln a)/da = 1/a
            self.grad += out.grad / self.data
        out._backward = _backward
        return out

    # --- everything else is built from the operations above --------------------
    def __neg__(self): return self * -1
    def __sub__(self, other): return self + (-other)
    def __truediv__(self, other): return self * other ** -1
    def __radd__(self, other): return self + other       # lets you write 2 + v
    def __rmul__(self, other): return self * other       # and 2 * v
    def __rsub__(self, other): return Value(other) - self
    def __rtruediv__(self, other): return Value(other) / self
    def sigmoid(self): return 1 / (1 + (-self).exp())

    # --- the backward pass -------------------------------------------------------
    def backward(self):
        """Fill in .grad for every Value this one depends on."""
        order, seen = [], set()

        def visit(v):                  # depth-first: parents before children
            if v not in seen:
                seen.add(v)
                for p in v._parents:
                    visit(p)
                order.append(v)
        visit(self)

        self.grad = 1.0                # dL/dL = 1
        for v in reversed(order):      # children before parents
            v._backward()


# === 1. One ride, by hand in Chapter 4, now automatic ===========================
w, b = Value(3.0), Value(5.0)
x, y = 5, 21                           # the 5 km ride that took 21 minutes
pred = w * x + b
e = pred - y
L = e ** 2
L.backward()
print("one ride:   pred", pred.data, " error", e.data, " loss", L.data)
print("            dL/dw =", w.grad, " dL/db =", b.grad)

# === 2. All four rides: the same numbers Chapter 4 derived ======================
rides_km = [2, 5, 8, 12]
rides_min = [11, 21, 28, 42]
w, b = Value(3.0), Value(5.0)
L = sum((w * xi + b - yi) ** 2 for xi, yi in zip(rides_km, rides_min)) / len(rides_km)
L.backward()
print("\nfour rides: loss", L.data, " dL/dw =", w.grad, " dL/db =", b.grad)

# === 3. Gradient check against central differences ==============================
def mse(wv, bv):
    return sum((wv * xi + bv - yi) ** 2 for xi, yi in zip(rides_km, rides_min)) / 4

h = 1e-5
num_w = (mse(3 + h, 5) - mse(3 - h, 5)) / (2 * h)
num_b = (mse(3, 5 + h) - mse(3, 5 - h)) / (2 * h)
print("numeric:    dL/dw =", round(num_w, 6), " dL/db =", round(num_b, 6))

# === 4. Why the gradients use += : a value used twice ===========================
a = Value(3.0)
f = a * a                              # f = a², so df/da should be 2a = 6
f.backward()
print("\na*a at a=3: grad", a.grad, "(a node with two paths to the output adds both)")

# === 5. The picnic neuron from Chapters 2 and 3, with cross-entropy =============
xs = [0.7, 0.2, 0.5]
ws = [Value(2.0), Value(-3.0), Value(-1.0)]
bias = Value(0.5)
z = sum((wi * xi for wi, xi in zip(ws, xs)), bias)
p = z.sigmoid()
label = 1                              # the picnic went well
loss = -(label * p.log() + (1 - label) * (1 - p).log())
loss.backward()
print("\npicnic:     z", round(z.data, 4), " p", round(p.data, 4), " loss", round(loss.data, 4))
print("            dloss/dz =", round(z.grad, 4), "  (p - y =", round(p.data - label, 4), ")")
print("            dloss/dw =", [round(wi.grad, 4) for wi in ws], " dloss/db =", round(bias.grad, 4))

# === 6. How big is the graph? ===================================================
def count(v, seen=None):
    seen = set() if seen is None else seen
    if v not in seen:
        seen.add(v)
        for q in v._parents:
            count(q, seen)
    return len(seen)
print("\nnodes in the picnic graph:", count(loss), "  in the four-ride graph:", count(L))