tinygrad/tinygrad

tinygrad – a tiny, hackable deep‑learning stack

What it is – tinygrad is an open‑source Python library that implements a complete end‑to‑end deep‑learning framework, but with a deliberately small code base. It provides:

A Tensor class with automatic differentiation (autograd).
An IR / compiler that fuses operations and lowers them to device kernels.
JIT / graph execution (via TinyJit).
A lightweight nn / optim / dataset layer so you can train real models.

It aims to sit between the full‑featured PyTorch and the ultra‑minimal micrograd, offering PyTorch‑like ergonomics while exposing the compiler internals for easy tinkering.

How it compares to the big players

Feature	PyTorch	JAX	tinygrad
Eager Tensor API	✅	–	✅
Autograd	✅	✅ (via JAXPR)	✅
Functional transforms (vmap, pmap)	–	✅	partial
Compiler & IR visible	–	✅	✅
Front‑end (nn, optim, datasets)	✅	–	✅
Multi‑device support	✅	✅	✅ (CPU, CUDA, Metal, OpenCL, AMD, NV, QCOM, WebGPU)

Quick‑start example

from tinygrad import Tensor, nn, Context

class LinearNet:
    def __init__(self):
        self.l1 = Tensor.kaiming_uniform(784, 128)
        self.l2 = Tensor.kaiming_uniform(128, 10)
    def __call__(self, x: Tensor) -> Tensor:
        return x.flatten(1).dot(self.l1).relu().dot(self.l2)

model = LinearNet()
optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)

x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])
with Context(TRAINING=1):
    for i in range(10):
        optim.zero_grad()
        loss = model(x).sparse_categorical_crossentropy(y).backward()
        optim.step()
        print(i, loss.item())

The repository ships a full MNIST example (examples/beautiful_mnist.py) that reaches ~98 % accuracy in about five seconds on a GPU.

Device support (accelerators)

tinygrad can run on many back‑ends with only ~25 primitive ops per device:

CPU
CUDA (NVIDIA GPUs)
Metal (Apple GPUs)
OpenCL
AMD, NV, QCOM (specialised runtimes)
WebGPU (browser / WASM) Run python3 -c "from tinygrad import Device; print(Device.DEFAULT)" to see which accelerator will be used by default.

Installation

# From source (recommended)
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

# Or directly from the repo
python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

The project also provides a .[testing] extra for running the test suite.

Documentation & community

Docs & quick‑start: https://docs.tinygrad.org/ (generated from the docs/ folder)
Discord: https://discord.gg/ZjZadyC7PK – active chat for users and developers.
Contributing: The README outlines a strict PR policy (no code‑golf, benchmark any speed claims, include regression tests, etc.) and mentions cash bounties for accepted improvements.

When to use tinygrad

Learning & experimentation – because the compiler pipeline is exposed, you can see exactly how high‑level tensor ops become low‑level kernels.
Prototyping on limited hardware – the library is lightweight and runs on CPUs, GPUs, and even WebGPU.
Research that needs a modifiable stack – you can add new ops or change scheduling with only a few files.

If you need the massive ecosystem, production‑grade stability, or the full suite of JAX transforms, PyTorch or JAX remain the better choices.

TL;DR

tinygrad is a real, functional deep‑learning framework that implements tensors, autograd, a JIT‑based compiler, and multi‑device execution in a compact, readable code base. It’s ideal for students, hobbyists, and researchers who want to understand or extend the internals of a modern ML stack without the overhead of a huge library.