Somewhere between PyTorch’s twenty million lines and micrograd’s two hundred lies the sweet spot for understanding how deep learning frameworks actually work. tinygrad, maintained by tiny corp and released under the MIT license, plants its flag exactly there: an end-to-end stack containing a tensor library with autograd, an intermediate representation and compiler that fuse and lower kernels, a JIT with graph execution, and enough nn, optim, and dataset machinery to train real models. The README is candid about its lineage — PyTorch for ergonomics, JAX for IR-based autodiff, TVM for scheduling and codegen — and about the design law behind it all: reduce complexity, delete code, and never confuse fewer lines with less clarity.
What makes the repository genuinely singular is where the ambition bottoms out. Most frameworks stop at the vendor’s runtime and call it portability. tinygrad ships userspace GPU drivers written from scratch for AMD and NVIDIA hardware, built on tens of thousands of autogenerated register definitions, so the path from a Python expression to commands on the silicon is entirely visible in one tree. Reading it end to end, you watch a lazy graph become a schedule, a schedule become a program, a program become machine code, and machine code become work on a queue — with nothing hidden behind a binary blob. As always in this series, what follows is an educational tour of published source code.
tinygrad overview: the Tensor facade builds a lazy UOp graph, the schedule and codegen passes lower it to programs, renderers print source or machine code, and the device layer launches everything through hand-written GPU drivers.
Reading the overview from left to right:
- the Tensor facade — 508 lines of PyTorch-shaped API where nothing computes until realize().
- the op mixins — elementwise, movement, reduce, and gradient behavior composed onto Tensor.
- the UOp IR — one universal graph node type plus a pattern-matching rewriter.
- the schedule passes — rewrites that turn tensor graphs into kernel-shaped form.
- the codegen — assembles programs, with BEAM search over scheduling choices.
- the renderers — print programs as CUDA, OpenCL, Metal, or WGSL source.
- the device layer — Buffer, Compiler, Program, and Compiled, the contract every backend fulfills.
- the AMD driver — a userspace GPU driver written from scratch on autogenerated register headers.
Why You Need This
The first reason is that tinygrad is the only project of its scale where you can hold the whole mental model in your head. The sz.py script in the repository root counts tokens per file and excludes the autogenerated and asset directories precisely because the hand-written core is meant to stay surveyable. The Tensor facade is a few hundred lines, the IR is one Python class with a pattern rewriter, and every stage between your loss function and your GPU is a named file you can read in one sitting. If you have ever bounced off PyTorch’s internals or XLA’s opacity, this is the codebase that finally shows the entire machine.
The second reason is the compiler education. The path from graph to kernel is the same one production systems use, just legible: schedule passes perform rangeification so tensor ops become explicit loops, indexing passes derive the address arithmetic, decomposition breaks transcendental functions into polynomial approximations of primitive ALU operations, and late passes coalesce and allocate registers before a renderer prints source for CUDA, OpenCL, Metal, WGSL, LLVM IR, or, for AMD, real machine code. On top of it all sits BEAM search, which compiles the same kernel many ways and keeps the fastest — the same idea TVM industrialized, shrunk to a search over a handful of well-chosen knobs. These are transferable concepts whether you ever touch GPU code or not.
The third reason is the drivers. ops_amd.py and ops_nv.py, at 861 and 673 lines, implement the backend contract directly against the hardware, supported by a hardware compute queue abstraction and autogenerated register definitions — the NVIDIA headers alone run past 26,000 generated lines across driver versions. Supporting a new accelerator only requires about 25 low-level operations, which is why the backend list already spans OpenCL, CPU, Metal, CUDA, AMD, NVIDIA, Qualcomm, and WebGPU. If you want to understand what a driver actually does — queues, signals, mappings, doorbells — without enrolling in a kernel course, these files are the friendliest introduction in existence.
The detailed view: the Tensor facade and its mixins on the left, the UOp IR core, the schedule and codegen stack, execution and backends on the right, with nn, ONNX, and the graph visualizer around the edges.
Start at the facade, because it explains the project’s central trick: laziness. tensor.py (508 lines) gives you the familiar API, but every operation returns more graph. The behavior lives in the mixin package — op.py (1,724 lines) is the base, with elementwise.py (877 lines), movement.py (578 lines), rand.py (276 lines), gradient.py (137 lines), and companions registering their operations. Everything those mixins build is a UOp, defined in uop/ops.py at line 230: one node type carries operations, dtypes, sources, and arguments for the entire stack. The same file holds the rewrite engine — UPat at line 1439, PatternMatcher at line 1582, graph_rewrite at line 1886 — so optimization is literally a list of pattern-to-replacement rules you can read like flashcards. Supporting cast in the uop package keeps the IR honest: symbolic.py (432 lines) simplifies symbolic expressions, spec.py (228 lines) validates graph well-formedness, and dtype.py (290 lines) defines the type lattice down to pointer and vector types.
The compile stack turns that graph into programs. schedule/init.py (256 lines) chains the lowering passes: rangeify.py (316 lines) converts tensor-level operations into explicit loop ranges, indexing.py (277 lines) derives the memory addressing math, and multi.py (275 lines) splits work across devices with allreduce coordination. codegen/init.py (435 lines) then assembles programs, calling on decomp/transcendental.py (260 lines) to reduce exp, log, and sin to primitive operations, opt/search.py (159 lines) to run BEAM search when enabled, and the late passes such as coalesce.py (152 lines) to finish the job. Renderers are the last stop before hardware: cstyle.py (556 lines) covers the C-family dialects in one printer, llvmir and wgsl handle their targets, and the AMD renderer generate.py (511 lines) emits actual RDNA and CDNA machine instructions — the same code paths the tinygrad-kernels experiments in llm/kernels build upon.
Execution threads it together. engine/realize.py (260 lines) exposes the compile, link, and run stages that fire when a Tensor’s realize() is called, the moment laziness ends. engine/jit.py (161 lines) implements TinyJit: on capture it substitutes parametric buffer slots into the recorded linear graph and fixes a memory plan, and on every later call it replays the compiled artifact with fresh inputs — jit_lower even honors a JITBEAM knob to re-search schedules at capture time. device.py (455 lines) defines the contract — Buffer at line 109, Allocator at 247, Compiler at 321, Program at 356, Compiled at 360 — and the backends implement it. The GPU backends go through support/hcq2.py (474 lines), a hardware compute queue layer of signals, submissions, and mappings that both the AMD and NVIDIA drivers share. The autogen directory, with machine-parsed register and instruction definitions for NVIDIA, AMD, Qualcomm, Apple, and more, is what makes hand-written drivers of this size reviewable at all.
Around the core sits a complete practical toolkit. nn/init.py (353 lines) supplies layers and the optim module (159 lines) behind the README’s famous example — beautiful_mnist.py reaches 98 percent accuracy in about five seconds on a GPU. nn/state.py (264 lines) handles checkpoint save and load, nn/onnx.py (1,058 lines) executes ONNX models as tinygrad graphs, and the llm package adds GGUF loading, a model runner, a CLI, and a server. For debugging and teaching there is viz/serve.py (688 lines), a local server that renders every rewrite pass of your graph, and helpers.py (585 lines) provides the ContextVar knobs — DEBUG levels that print the fused kernels, BEAM, CAPTURE_PROCESS_REPLAY — that make the whole stack observable. The repository even monitors itself: process replay tests compare the kernels your change generates against master, so refactors cannot silently alter codegen.
From Install to a Fused Matmul
The recommended install is from source, and the first demo is a single fused kernel:
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .
python3 -c "from tinygrad import Device; print(Device.DEFAULT)"
DEBUG=3 python3 -c "from tinygrad import Tensor; N=1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N); (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"
The DEBUG=3 run prints the kernels the compiler chose; change it to DEBUG=4 to see the generated code itself. From there the README’s neural network example is a complete training loop in a dozen lines — Tensor.kaiming_uniform weights, nn.optim.Adam, a backward call, a Context(TRAINING=1) guard — and examples/beautiful_mnist.py is the fuller version worth running. Tests live under the test directory, the contributing guide asks for regression tests with every feature, and the project pays cash bounties for improvements, with a strict no-AI-written-PRs policy that tells you exactly how seriously the maintainers take human readability. In a field where that is rare, the code itself is the best evidence it works.
Honest limits: tinygrad is explicitly pre-1.0, so APIs move and documentation, while improving, still trails the code; the functional-transform surface is smaller than JAX’s, with no full vmap or pmap equivalents yet; debugging miscompiled kernels demands compiler literacy that eager frameworks let you avoid; and the driver work, while thrilling, targets a specific set of silicon rather than whatever CUDA supports. None of that dims the achievement: this is the most complete, most readable demonstration in open source that a deep learning framework is not magic — it is a graph, a list of rewrite rules, a code generator, and a driver, all of which you can read in a week. Enjoyed this post? Never miss out on future posts by following us