Building an Autograd Engine from Scratch in 200 Lines of Pure Python

1. What is Automatic Differentiation?

When developers first train a model in PyTorch or JAX, they invoke:

loss.backward()
optimizer.step()

Behind this concise API lies an elegant data structure: a dynamically traced Directed Acyclic Graph (DAG) of mathematical operations.

Automatic differentiation is neither numerical differentiation (finite differences with approximation errors) nor symbolic differentiation (which suffers from expression swell). Instead, it evaluates exact analytic derivatives at runtime by accumulating local Jacobians using the Multivariate Chain Rule.

2. The Value Node & Computational Graph

To build an autograd engine, we encapsulate every scalar number into a Value class that remembers:

1. data: The current numerical scalar.
2. grad: The accumulated derivative $\frac{\partial \mathcal{L}}{\partial \text{self}}$.
3. _prev: Pointers to its parent operand nodes.
4. _backward: A closure that distributes downstream incoming gradient to parents.

class Value:
    def __init__(self, data, _children=(), _op=""):
        self.data = float(data)
        self.grad = 0.0
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op

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

3. Operator Overloading with Local Backward Hooks

Let's implement addition and multiplication.

Addition: $z = x + y$

By the chain rule, $\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial z} \cdot 1$. The gradient simply splits equally:

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 += 1.0 * out.grad
        other.grad += 1.0 * out.grad
    out._backward = _backward

    return out

[!IMPORTANT] Notice the += accumulation! If a variable is reused in multiple operations (e.g. $y = x + x$), its total derivative is the sum of gradients across all paths. Using = instead of += is the #1 bug in custom autodiff engines!

Multiplication: $z = x \cdot y$

By the product rule: $\frac{\partial \mathcal{L}}{\partial x} = y \cdot \frac{\partial \mathcal{L}}{\partial z}, \quad \frac{\partial \mathcal{L}}{\partial y} = x \cdot \frac{\partial \mathcal{L}}{\partial z}$

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
        other.grad += self.data * out.grad
    out._backward = _backward

    return out

Non-Linear Activation (Tanh)

$f(x) = \tanh(x), \quad \frac{df}{dx} = 1 - \tanh^2(x)$

import math

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

    def _backward():
        self.grad += (1.0 - t**2) * out.grad
    out._backward = _backward

    return out

4. Topological Sorting of the DAG

When backpropagating through an arbitrary graph, a parent node must not execute its _backward() until all of its downstream children have completely finished sending their gradients.

To guarantee correct execution order, we construct a Topological Sort using Depth-First Search (DFS):

def backward(self):
    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)

    # Base case: seed the output loss with dL/dL = 1.0
    self.grad = 1.0

    # Execute in reverse topological order (leaves to inputs)
    for node in reversed(topo):
        node._backward()

5. Training a Neuron with Zero External Libraries

Now let's wire these primitive scalar nodes into a trainable artificial neuron:

# Inputs
x1 = Value(2.0)
x2 = Value(0.0)

# Trainable Weights and Bias
w1 = Value(-3.0)
w2 = Value(1.0)
b = Value(6.88137)

# Forward pass: y = tanh(x1*w1 + x2*w2 + b)
n = (x1 * w1) + (x2 * w2) + b
out = n.tanh()

# Backward pass
out.backward()

print(f"Output activation: {out.data:.4f}")
print(f"dL/dw1: {w1.grad:.4f}")
print(f"dL/dw2: {w2.grad:.4f}")

This output gives exact analytical gradients identical to PyTorch's torch.autograd. Understanding this at the bytecode level eliminates any fear of backprop.