tinygrad

tinygrad: For something between PyTorch and karpathy/micrograd. Maintained by tiny corp. Homepage | Documentation | Discord tinygrad is an end-to-end deep learning stack: Tensor library with autograd IR and compiler that fuse and lower kernels JIT + graph execution nn / optim / datasets for real training It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable. How tinygrad compares PyTorch ✅ Similar: eager Tensor API, autograd, optim, basic datasets and layers. ✅ You can write familiar training loops. 🔁 Unlike PyTorch, the entire compiler and IR are visible and hackable. JAX ✅ IR-based autodiff over primitives (like JAXPR + XLA). ✅ Function-level JIT (TinyJit) that captures and replays kernels. 🔁 Fewer functional transforms (no full vmap/pmap yet), but far easier to read. TVM ✅ Multiple lowering passes, scheduling, and BEAM search over kernels. ✅ Device “graphs” for batched execution. 🔁 tinygrad also ships the front-end framework (tensors, nn, optim), not just the compiler. Laziness Try a matmul. See how, despite the style, it is fused into one kernel with the power of laziness. 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()" And we can change DEBUG to 4 to see the generated code. Neural networks As it turns out, 90% of what you need for neural networks are a decent autograd/tensor library. Throw in an optimizer, a data loader, and some compute, and you have all you need. 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]) # replace with real mnist dataloader 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()) See examples/beautiful_mnist.py for the full version that gets 98% in ~5 seconds Accelerators tinygrad already supports numerous accelerators, including: OpenCL CPU METAL CUDA AMD NV QCOM WEBGPU And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops. To check default accelerator run: python3 -c "from tinygrad import Device; print(Device.DEFAULT)" Installation The current recommended way to install tinygrad is from source. From source git clone https://github.com/tinygrad/tinygrad.git cd tinygrad python3 -m pip install -e . Direct (master) python3 -m pip install git+https://github.com/tinygrad/tinygrad.git Documentation Documentation along with a quick start guide can be found on the docs website built from the docs/ directory. Quick example comparing to PyTorch from tinygrad import Tensor x = Tensor.eye(3).clone() # clone to make it a buffer y = Tensor([[2.0,0,-2.0]]) z = y.matmul(x).sum() z.backward() print(x.grad.tolist()) # dz/dx print(y.grad.tolist()) # dz/dy The same thing but in PyTorch: import torch x = torch.eye(3, requires_grad=True) y = torch.tensor([[2.0,0,-2.0]], requires_grad=True) z = y.matmul(x).sum() z.backward() print(x.grad.tolist()) # dz/dx print(y.grad.tolist()) # dz/dy Contributing There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted. If you do submit a PR, please include a sentence or two about why you want this merged and why you think it will improve the project. If you are a new contributor with something that looks even close to AI written, it will be closed without feedback and you may be banned from our GitHub. No human should waste time reading AI slop. And for everyone, if you used AI, disclose what you used it for. We'll start with what will get your PR closed with a pointer to this section: No code golf! While low line count is a guiding light of this project, anything that remotely looks like code golf will be closed. The true goal is reducing complexity and increasing readability, and deleting \ns does nothing to help with that. All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless and carry a risk of introducing bugs. Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainability and readability. In general, the code outside the core tinygrad/ folder is not well tested, so unless the current code there is broken, you shouldn't be changing it. If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review. Now, what we want: Bug fixes (with a regression test) are great! This library isn't 1.0 yet, so if you stumble upon a bug, fix it, write a test, and submit a PR, this is valuable work. Solving bounties! tinygrad offers cash bounties for certain improvements to the library. All new code should be high quality and well tested. Features. However, if you are adding a feature, consider the line tradeoff. If it's 3 lines, there's less of a bar of usefulness it has to meet over something that's 30 or 300 lines. All features must have regression t