chainrule

A scalar autograd engine for learning backpropagation: a Value class, reverse-mode backward(), a tiny nn module, and two things most micrograd-style toy engines skip — a way to actually look at the graph after backward() runs, and a way to check that your gradients are right before you trust them.

$ chainrule demo
chainrule demo -- XOR, 200 steps, 17 parameters, trained in 79.4 ms
loss  step 0    = 1.0419
loss  step 199  = 0.0035
loss  final     = 0.0034

Why this exists

Backpropagation is one of those things that looks obvious once you have implemented it and completely opaque before. Reading someone else's _backward closure is not the same as writing your own and then finding out, from a wrong gradient, exactly which line was off by a sign. This package is the small version of that exercise: a scalar Value type whose forward and backward passes fit on one screen, plus the two tools you actually want while you are debugging one — a picture of the graph, and a numerical check on the gradient it produced.

Install

$ pip install chainrule

Or, if you use uv:

$ uv tool install chainrule

Pure standard library plus typer and rich for the CLI. No numeric dependencies — every gradient here is one Python float.

The graph, on a classic example

from chainrule import Value

a = Value(2.0)
b = Value(-3.0)
c = Value(10.0)
f = Value(-2.0)

L = (a * b + c) * f
L.backward()
print(L.graph())

#0 data=-8.0000 grad=+1.0000 [*]
   ├─ #1 data=+4.0000 grad=-2.0000 [+]
   │  ├─ #2 data=-6.0000 grad=-2.0000 [*]
   │  │  ├─ #3 data=+2.0000 grad=+6.0000 [leaf]
   │  │  └─ #4 data=-3.0000 grad=-4.0000 [leaf]
   │  └─ #5 data=+10.0000 grad=-2.0000 [leaf]
   └─ #6 data=-2.0000 grad=+4.0000 [leaf]

Every node shows its forward value and, after backward(), the gradient that flowed into it. A node used more than once anywhere in the graph — a weight shared by two neurons, an input fed to a whole layer — is only expanded the first time; later occurrences print -> see #N instead of redrawing the whole subtree, so the tree stays readable even for a small MLP. Pass graph(max_depth=...) to cut off a bigger one.

Checking a gradient before you trust it

from chainrule import Value, gradcheck

def f(xs):
    x, y = xs
    return x.tanh() * y + (x * y).exp()

result = gradcheck(f, [0.7, -0.3])
print(result.max_rel_error, result.passed)

gradcheck runs f once to get the analytic gradient from backward(), then re-runs it with each input nudged by +-eps to get a central finite difference, and reports the relative error between the two. A correct _backward should land many orders of magnitude under the default tol=1e-4; a wrong sign or a missing term usually shows up as an error close to 1.

A tiny neural net

from chainrule import MLP, Value

model = MLP(2, [4, 1], activation="tanh", out_activation="tanh")
pred = model([Value(1.0), Value(-1.0)])
loss = (pred - 1.0) ** 2
loss.backward()

for p in model.parameters():
    p.data -= 0.1 * p.grad

Neuron, Layer and MLP are plain Python loops over Value objects — no batching, no tensors — so the forward pass you read in nn.py is exactly the graph Value.graph() will show you afterwards. softmax_cross_entropy(logits, target) and softmax(logits) use the log-sum-exp shift, so they stay finite for logits in the thousands instead of overflowing math.exp.

CLI

$ chainrule demo          # trains a 2-4-1 MLP on XOR and prints the final graph
$ chainrule gradcheck      # runs the built-in gradient checks and exits 1 on a failure

$ chainrule gradcheck
check        max rel. error  status
polynomial         2.43e-10  pass
activations        2.27e-11  pass
division           1.44e-10  pass
tiny_mlp           6.17e-11  pass

When not to use this

For anything beyond a learning exercise, use a real framework — PyTorch if you want to stay in Python, or JAX. This engine differentiates one float at a time with no batching, no GPU support, and no numerical stability work beyond sigmoid and the log-sum-exp trick in softmax. It is fast enough to train the XOR demo in under a second and not fast enough for anything that needs a real tensor library.

License

MIT. See LICENSE.