tinygrad/tinygrad

tinygrad – 一个精简且可修改的深度学习堆叠

它是什么 – tinygrad 是一个开源 Python 库，实现了一个完整的端到端深度学习框架，但代码库刻意保持精简。它提供：

具有自动微分 (autograd) 的 Tensor 类。
一个将运算融合并降低为设备核心 (device kernels) 的 IR / 编译器。
JIT / 图执行 (通过 TinyJit)。
轻量级的 nn / optim / dataset 层，让您可以训练真实模型。

它的目标是介于功能齐全的 PyTorch 与极简的 micrograd 之间，在提供类似 PyTorch 的人体工学体验的同时，公开编译器内部结构以便于调整。

与大型框架的比较

特性	PyTorch	JAX	tinygrad
Eager Tensor API	✅	–	✅
Autograd	✅	✅ (via JAXPR)	✅
函数式转换 (vmap, pmap)	–	✅	部分
编译器与 IR 可见	–	✅	✅
前端 (nn, optim, datasets)	✅	–	✅
多设备支持	✅	✅	✅ (CPU, CUDA, Metal, OpenCL, AMD, NV, QCOM, WebGPU)

快速开始示例

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())

该存储库附带一个完整的 MNIST 示例 (examples/beautiful_mnist.py)，在 GPU 上约五秒内即可达到约 98% 的准确率。

设备支持 (加速器)

tinygrad 可以在许多后端上执行，每个设备仅需约 25 个原始运算：

CPU
CUDA (NVIDIA GPU)
Metal (Apple GPU)
OpenCL
AMD, NV, QCOM (专用执行环境)
WebGPU (浏览器 / WASM) 执行 python3 -c "from tinygrad import Device; print(Device.DEFAULT)" 以查看默认使用的加速器。

安装

# 从源码安装 (推荐)
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

# 或直接从存储库安装
python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

该项目还提供 .[testing] 选项以执行测试套件。

文档与社区

文档与快速开始: https://docs.tinygrad.org/ (由 docs/ 文件夹生成)
Discord: https://discord.gg/ZjZadyC7PK – 用户与开发者的活跃聊天室。
贡献: README 概述了严格的 PR 政策（无代码高尔夫、对任何速度声明进行基准测试、包含回归测试等），并提到对被采纳的改进提供现金奖励。

何时使用 tinygrad

学习与实验 – 由于编译器管线是公开的，您可以精确地观察高阶张量运算如何转变为低阶核心。
在受限硬件上进行原型设计 – 该库轻量且可在 CPU、GPU 甚至 WebGPU 上执行。
需要可修改堆叠的研究 – 您只需修改几个文件即可新增运算或变更调度。

如果您需要庞大的生态系统、生产级的稳定性或完整的 JAX 转换套件，PyTorch 或 JAX 仍然是更好的选择。

总结

tinygrad 是一个真实、功能齐全的深度学习框架，在精简且易读的代码库中实现了张量、自动微分、基于 JIT 的编译器以及多设备执行。它非常适合想要理解或扩展现代 ML 堆叠内部结构，且不想承担庞大库负担的学生、爱好者与研究人员。