Package: tinygrad
tinygrad: For something betweenPyTorchandkarpathy/micrograd. Maintained bytiny corp.Homepage|Documentation|Examples|Showcase|DiscordThis 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 weraised some moneyto make it good. Someday, we will tape out chips.FeaturesLLaMA and Stable Diffusiontinygrad can runLLaMAandStable Diffusion!LazinessTry a matmul. See how, despite the style, it is fused into one kernel with the power of laziness.DEBUG=3python3-c"from tinygrad import Tensor;N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2);print((c.numpy() - (a.numpy() @ b.numpy())).mean())"And we can changeDEBUGto4to see the generated code.Neural networksAs 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.fromtinygradimportTensor,nnclassLinearNet:def__init__(self):self.l1=Tensor.kaiming_uniform(784,128)self.l2=Tensor.kaiming_uniform(128,10)def__call__(self,x:Tensor)->Tensor:returnx.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 dataloaderforiinrange(10):optim.zero_grad()loss=model(x).sparse_categorical_crossentropy(y).backward()optim.step()print(i,loss.item())Seeexamples/beautiful_mnist.pyfor the full version that gets 98% in ~5 secondsAcceleratorstinygrad already supports numerous accelerators, including:CPUGPU (OpenCL)C Code (Clang)LLVMMETALCUDAPyTorchHIPAnd it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
More information can be found in thedocumentation for adding new accelerators.InstallationThe current recommended way to install tinygrad is from source.From sourcegitclonehttps://github.com/tinygrad/tinygrad.gitcdtinygrad
python3-mpipinstall-e.Direct (master)python3-mpipinstallgit+https://github.com/tinygrad/tinygrad.gitDocumentationDocumentation along with a quick start guide can be found in thedocs/directory.Quick example comparing to PyTorchfromtinygradimportTensorx=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/dxprint(y.grad.numpy())# dz/dyThe same thing but in PyTorch:importtorchx=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/dxprint(y.grad.numpy())# dz/dyContributingThere has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted.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 uselessandcarry 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 maintainablity and readablity.In general, the code outside the coretinygrad/folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.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! tinygradoffers cash bountiesfor 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 tests. In general with no other constraints, your feature's API should match torch or numpy.Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win.Tests/fuzzers. If you can add tests that are non brittle, they are welcome. We have some fuzzers in here too, and there's a plethora of bugs that can be found with them and by improving them. Finding bugs, even writing broken tests (that should pass) with@unittest.expectedFailureis great. This is how we make progress.Dead code removal from coretinygrad/folder. We don't care about the code in extra, but removing dead code from the core library is great. Less for new people to read and be confused by.Running testsYou should install the pre-commit hooks withpre-commit install. This will run the linter, mypy, and a subset of the tests on every commit.For more examples on how to run the full test suite please refer to theCI workflow.Some examples of running tests locally:python3-mpipinstall-e'.[testing]'# install extra deps for testingpython3test/test_ops.py# just the ops testspython3-mpytesttest/# whole test suite