tinygrad

tinygrad sits somewhere between PyTorch and Karpathy's micrograd: a full deep-learning stack small enough to read end to end. You get an eager Tensor API with autograd, nn/optim/datasets for real training, and then the unusual part, a visible, hackable compiler: operations are lazy, fused into single kernels, lowered through multiple passes with scheduling and BEAM search over kernel choices, and replayed with TinyJit. DEBUG=3 shows you the fused kernel; DEBUG=4 shows the generated code.

It borrows ergonomics from PyTorch, IR-based autodiff ideas from JAX, and scheduling from TVM, but keeps everything in one small codebase maintained by tiny corp (George Hotz's company). At 33,700+ stars it is both a serious research vehicle and the best way to actually understand how a modern ML compiler works. The honest trade-off: it is intentionally minimal, so you will miss PyTorch's library breadth, and some backends and ops are community-grade rather than production-hardened.

1. Install: pip install tinygrad (Python 3.9+).
2. Write a familiar training loop with Tensor, nn, and tinygrad.nn.optim.
3. Set DEBUG=3 to see kernel fusion in action, DEBUG=4 to inspect generated code.
4. Wrap hot loops in TinyJit for capture-and-replay execution.
5. Target different accelerators via the documented backend env vars (see docs.tinygrad.org).