Build an Autograd Engine
Goal: Build a tiny automatic differentiation engine from scratch — a Value class that tracks computation graphs and computes gradients. Then build a neural network on top of it. Inspired by Karpathy’s micrograd.
Prerequisites: Backpropagation, Chain Rule, Derivatives, 07 - Backpropagation Step by Step
Why Autograd?
Tutorial 05 computed gradients by hand with matrix math. Tutorial 07 traced them through a tiny network with specific numbers. Both are rigid — change the architecture and you rewrite the math.
Autograd is the real solution: build a computation graph as you compute, then walk it backward to get all gradients automatically. This is what PyTorch does under the hood.
The Value Class
Every number becomes a Value that remembers how it was created:
import math
import numpy as np
import matplotlib.pyplot as plt
class Value:
def __init__(self, data, _children=(), _op=''):
self.data = data
self.grad = 0.0
self._backward = lambda: None # function that computes local gradients
self._prev = set(_children)
self._op = _op
def __repr__(self):
return f"Value(data={self.data:.4f}, grad={self.grad:.4f})"
def __add__(self, other):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data + other.data, (self, other), '+')
def _backward():
self.grad += out.grad # d(a+b)/da = 1
other.grad += out.grad # d(a+b)/db = 1
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():
self.grad += other.data * out.grad # d(a*b)/da = b
other.grad += self.data * out.grad # d(a*b)/db = a
out._backward = _backward
return out
def __pow__(self, other):
assert isinstance(other, (int, float))
out = Value(self.data ** other, (self,), f'**{other}')
def _backward():
self.grad += other * (self.data ** (other - 1)) * out.grad
out._backward = _backward
return out
def tanh(self):
t = math.tanh(self.data)
out = Value(t, (self,), 'tanh')
def _backward():
self.grad += (1 - t**2) * out.grad
out._backward = _backward
return out
def relu(self):
out = Value(max(0, self.data), (self,), 'relu')
def _backward():
self.grad += (self.data > 0) * out.grad
out._backward = _backward
return out
def exp(self):
e = math.exp(self.data)
out = Value(e, (self,), 'exp')
def _backward():
self.grad += e * out.grad
out._backward = _backward
return out
def backward(self):
"""Backprop through the entire computation graph."""
# Topological sort — process children before parents
topo = []
visited = set()
def build_topo(v):
if v not in visited:
visited.add(v)
for child in v._prev:
build_topo(child)
topo.append(v)
build_topo(self)
self.grad = 1.0 # dL/dL = 1
for v in reversed(topo):
v._backward()
# Make operations work both ways: 2 * Value and Value * 2
def __radd__(self, other): return self + other
def __rmul__(self, other): return self * other
def __neg__(self): return self * -1
def __sub__(self, other): return self + (-other)
def __rsub__(self, other): return (-self) + other
def __truediv__(self, other): return self * (other ** -1) if isinstance(other, Value) else self * Value(other) ** -1Test: Manual Computation
# f(a, b) = a * b + b^2
a = Value(2.0)
b = Value(3.0)
c = a * b # 6.0
d = b ** 2 # 9.0
e = c + d # 15.0
e.backward()
print(f"e = {e}") # 15.0
print(f"de/da = {a.grad}") # b = 3.0
print(f"de/db = {b.grad}") # a + 2b = 2 + 6 = 8.0Verify with PyTorch
import torch
a_t = torch.tensor(2.0, requires_grad=True)
b_t = torch.tensor(3.0, requires_grad=True)
e_t = a_t * b_t + b_t ** 2
e_t.backward()
print(f"PyTorch: de/da = {a_t.grad.item()}, de/db = {b_t.grad.item()}")The Key Insight: grad +=
Why += and not =? Because a value might be used multiple times:
a = Value(3.0)
b = a + a # a is used twice
b.backward()
print(f"db/da = {a.grad}") # 2.0, not 1.0 — both paths contributeThis is the multivariate chain rule in action. Every path through the graph contributes to the gradient.
Build a Neural Network on Top
import random
class Neuron:
def __init__(self, n_inputs):
self.w = [Value(random.uniform(-1, 1)) for _ in range(n_inputs)]
self.b = Value(0.0)
def __call__(self, x):
act = sum((wi * xi for wi, xi in zip(self.w, x)), self.b)
return act.tanh()
def parameters(self):
return self.w + [self.b]
class Layer:
def __init__(self, n_in, n_out):
self.neurons = [Neuron(n_in) for _ in range(n_out)]
def __call__(self, x):
out = [n(x) for n in self.neurons]
return out[0] if len(out) == 1 else out
def parameters(self):
return [p for n in self.neurons for p in n.parameters()]
class MLP:
def __init__(self, n_in, layer_sizes):
sizes = [n_in] + layer_sizes
self.layers = [Layer(sizes[i], sizes[i+1]) for i in range(len(layer_sizes))]
def __call__(self, x):
for layer in self.layers:
x = layer(x)
return x
def parameters(self):
return [p for layer in self.layers for p in layer.parameters()]Train on a Toy Dataset
# Binary classification: 4 points
X = [
[2.0, 3.0, -1.0],
[3.0, -1.0, 0.5],
[0.5, 1.0, 1.0],
[1.0, 1.0, -1.0],
]
y_true = [1.0, -1.0, -1.0, 1.0] # tanh outputs in [-1, 1]
model = MLP(3, [4, 4, 1]) # 3 inputs → 4 → 4 → 1 output
print(f"Number of parameters: {len(model.parameters())}")
# Training loop
losses = []
for epoch in range(100):
# Forward pass
preds = [model(x) for x in X]
loss = sum((p - yt) ** 2 for p, yt in zip(preds, y_true))
losses.append(loss.data)
# Backward pass
# Zero gradients first!
for p in model.parameters():
p.grad = 0.0
loss.backward()
# Update
lr = 0.05
for p in model.parameters():
p.data -= lr * p.grad
if epoch % 20 == 0:
print(f"Epoch {epoch:3d} | Loss: {loss.data:.4f} | Preds: {[f'{p.data:.2f}' for p in preds]}")
plt.plot(losses)
plt.xlabel("Epoch"); plt.ylabel("Loss")
plt.title("Autograd engine training")
plt.show()Visualize the Computation Graph
def draw_graph(root):
"""Print computation graph as text tree."""
def _draw(v, prefix="", is_last=True):
connector = "└── " if is_last else "├── "
print(f"{prefix}{connector}{v._op or 'input'} → {v.data:.4f} (grad={v.grad:.4f})")
children = list(v._prev)
for i, child in enumerate(children):
extension = " " if is_last else "│ "
_draw(child, prefix + extension, i == len(children) - 1)
_draw(root)
# Small example
a = Value(2.0); b = Value(-3.0); c = Value(10.0)
d = a * b + c
d.backward()
draw_graph(d)Why Zero Gradients?
The += in backward means gradients accumulate. If you don’t zero them before each backward pass:
a = Value(3.0)
b = a * 2
b.backward()
print(f"After 1st backward: {a.grad}") # 2.0
b = a * 2
b.backward()
print(f"After 2nd backward: {a.grad}") # 4.0 — accumulated!This is why PyTorch needs optimizer.zero_grad() before every step.
Exercises
-
Add more ops: Implement
log(),sigmoid(), and__matmul__on Value. The sigmoid backward is . -
Softmax + cross-entropy: Implement a multi-class classifier. You’ll need
exp()and normalization. Make a 2D spiral dataset and classify it. -
Numerical gradient check: For every operation, verify the analytical gradient matches .
-
Trace PyTorch: Look at
torch.Tensor— it has.grad_fn,.requires_grad, and.backward(). Now you know exactly what they do. -
Performance: This engine is ~1000x slower than PyTorch because it operates on scalars, not tensors. Think about how you’d extend
Valueto hold arrays while keeping the same API.
Next: 16 - Bigram Language Model — use this same gradient machinery to model language.