micrograd-rust

I originally wanted to implement micrograd in C and learn Rust separately; when the time came I thought why not just do both at once and write it in Rust instead. Was surprisingly some of the most fun I've had in a long time.

Autograd engine and small neural net library in Python and Rust, based on Karpathy's micrograd. Also available as a Python package; pip install micrograd-rust and you get the same API as micrograd but backed by Rust.

Install as a Python package

pip install micrograd-rust

from micrograd_rust import Value, MLP

a = Value(2.0)
b = Value(3.0)
c = (a * b).tanh()
c.backward()
print(a.grad)

Uses PyO3 to call into Rust under the hood so you get the speed without leaving Python.

What it does

So basically you build up a computation graph out of scalar values; every time you do like a + b or a * tanh(b) it creates a node that remembers what operation produced it and what the inputs were. Then you call backward() on the output and it walks the graph in reverse topological order, applying the chain rule at every node to compute all the gradients. That's it; that's all of autograd. No tensors, no batching, just individual floats with grad tracking; but it's enough to train a small MLP from scratch with no frameworks which is super cool.

Benchmarks

50 forward+backward passes per network size. Rust compiled with --release. PyTorch included for context.

Network	Params	Python	Rust	PyTorch	Rust vs Python
3->4->4->1	41	0.13ms	0.01ms	0.08ms	13x
3->16->16->1	353	1.32ms	0.12ms	0.07ms	11x
3->32->32->1	1,217	5.95ms	0.38ms	0.09ms	16x
3->64->64->1	4,481	30.41ms	1.18ms	0.08ms	26x
3->128->128->1	17,153	271.22ms	4.51ms	0.10ms	60x
3->256->256->1	67,073	1,496.61ms	19.96ms	0.12ms	75x
3->512->512->1	265,217	stack overflow	89.13ms	0.28ms	-
3->1024->1024->1	1,054,721	stack overflow	389.09ms	0.86ms	-

What's going on here?

Python is doing scalar operations one at a time, each with the overhead of Python's interpreter. Every single + goes through __add__, allocates a new Value object, appends to the graph, etc. At ~67K parameters the recursive topological sort blows the call stack entirely.

Rust does the same scalar operations but without interpreter overhead; direct machine code, no garbage collector, no dynamic dispatch. The speedup grows with network size (between 1-2 OOMs) because Python's per-operation cost is constant, and more operations means more accumulated overhead. Look at how it scales though; 13x at 41 params, 75x at 67K params; the ratio just keeps climbing because Rust's cost per op is basically nothing so you're really just measuring how much Python's overhead compounds.

PyTorch is in a different league because it's not actually doing scalar ops; under the hood it's dispatching to optimized C/BLAS/LAPACK routines that do batched matrix multiplications. It barely blinks going from 41 to 67K parameters; going from 67K to 1M params costs it about 0.7ms while Rust goes from 20ms to 389ms and Python is already dead. This is why real ML frameworks exist LOL!

Getting started (from source)

See python/README.md and rust/README.md for setup and usage.

Acknowledgments

Built following Karpathy's spelled out intro to neural nets video where he builds micrograd from scratch and explains backpropagation super clearly.