This may not be the best deep learning framework, but it is a deep learning framework.
Due to its extreme simplicity, it aims to be the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.
tinygrad is still alpha software, but we raised some money to make it good. Someday, we will tape out chips.
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.tensor import Tensor; N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N); c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2); print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
And we can change
4 to see the generated code.
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.tensor import Tensor import tinygrad.nn.optim as optim class TinyBobNet: def __init__(self): self.l1 = Tensor.uniform(784, 128) self.l2 = Tensor.uniform(128, 10) def forward(self, x): return x.dot(self.l1).relu().dot(self.l2).log_softmax() model = TinyBobNet() optim = optim.SGD([model.l1, model.l2], lr=0.001) # ... complete data loader here out = model.forward(x) loss = out.mul(y).mean() optim.zero_grad() loss.backward() optim.step()
tinygrad already supports numerous accelerators, including:
And it is easy to add more! Your accelerator of choice only needs to support a total of 26 (optionally 27) low level ops. More information can be found in the documentation for adding new accelerators.
The current recommended way to install tinygrad is from source.
git clone https://github.com/geohot/tinygrad.git cd tinygrad python3 -m pip install -e .
Don't forget the
. at the end!
Documentation along with a quick start guide can be found in the docs/ directory.
from tinygrad.tensor import Tensor x = Tensor.eye(3, requires_grad=True) y = Tensor([[2.0,0,-2.0]], requires_grad=True) z = y.matmul(x).sum() z.backward() print(x.grad.numpy()) # dz/dx print(y.grad.numpy()) # 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.numpy()) # dz/dx print(y.grad.numpy()) # dz/dy
There has been a lot of interest in tinygrad lately. Here are some basic guidelines for contributing:
Additional guidelines can be found in CONTRIBUTING.md.
For more examples on how to run the full test suite please refer to the CI workflow.
python3 -m pip install -e '.[testing]' python3 -m pytest python3 -m pytest -v -k TestTrain python3 ./test/models/test_train.py TestTrain.test_efficientnet