PyTorch ergonomics with a compiler you can actually read

What it does

tinygrad is an end-to-end deep learning framework with tensors, autograd, optimizers, and standard layers. It compiles and fuses kernels through an exposed IR and compiler, then executes them via JIT or graph mode across a wide range of accelerators. You write familiar eager-mode training loops, but the framework lazily fuses operations and lowers them to device-specific code behind the scenes.

The interesting bit

The project borrows PyTorch’s ergonomics, JAX’s IR-based autodiff, and TVM’s scheduling and codegen, yet keeps the entire compiler readable enough to modify. It achieves kernel fusion through laziness rather than explicit graph construction, and adding a new accelerator requires implementing only about 25 low-level ops. The maintainers also offer cash bounties for improvements and will close PRs that sacrifice readability for marginal speedups.

Key highlights

• End-to-end stack: tensors, autograd, nn/optim, compiler, JIT, and graph execution
• Lazy evaluation automatically fuses operations into single kernels
• Supports CPU, CUDA, Metal, AMD, NV, Qualcomm, WebGPU, and OpenCL backends
• Entire compiler and IR are exposed and hackable, unlike PyTorch’s black box
• New accelerators require only ~25 low-level ops to add

Caveats

• Not yet at version 1.0, so expect rough edges
• Lacks full JAX-style functional transforms such as vmap and pmap
• Contributing guidelines are strict: speedups must be benchmarked, and large or “complex” PRs are unlikely to be merged

Verdict

Worth a look if you want to understand how a modern deep learning compiler works without wading through millions of lines of C++. Skip it if you need production stability or JAX’s full suite of functional transforms today.