---

tinygrad 是一个端到端的深度学习栈：

• 张量库，支持自动微分
• 中间表示与编译器，用于融合和降低计算核
• 即时编译 + 图执行
• nn / optim / datasets，用于实际训练

它受到 PyTorch（易用性）、JAX（函数式变换和基于 IR 的自动微分）以及 TVM（调度和代码生成）的启发，但刻意保持轻量级且易于扩展。

---

tinygrad 与其他框架的对比

PyTorch

• ✅ 相似：急切模式下的 Tensor API、自动微分、优化器、基础数据集和层。
• ✅ 可以编写熟悉的训练循环。
• 🔁 不同于 PyTorch，tinygrad 的整个编译器和中间表示都是可见且可修改的。

JAX

• ✅ 基于中间表示的自动微分，使用原语（如 JAXPR + XLA）。
• ✅ 函数级别的 JIT（TinyJit），能够捕获并重放计算核。
• 🔁 功能性变换较少（尚未实现完整的 vmap/pmap），但代码更易读。

TVM

• ✅ 多次降低 pass、调度以及基于 BEAM 搜索的计算核选择。
• ✅ 设备“图”用于批处理执行。
• 🔁 tinygrad 同时提供了前端框架（张量、神经网络、优化器），而不仅仅是编译器。

---

惰性求值

尝试一个矩阵乘法。你会发现，尽管风格不同，但由于惰性求值机制，它会被融合为一个单一的计算核。

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

将 DEBUG 改为 4，即可查看生成的代码。

神经网络

事实证明，构建神经网络所需的 90% 仅仅是良好的自动微分/张量库。再加上优化器、数据加载器和足够的计算资源，你就拥有了所有必需的工具。

from tinygrad import Tensor, nn

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])  # 替换为真实的 MNIST 数据加载器

with Tensor.train():
  for i in range(10):
    optim.zero_grad()
    loss = model(x).sparse_categorical_crossentropy(y).backward()
    optim.step()
    print(i, loss.item())

完整版本请参见 examples/beautiful_mnist.py，该示例可在约 5 秒内达到 98% 的准确率。

加速器支持

tinygrad 已经支持多种加速器，包括：

• OpenCL
• CPU
• METAL
• CUDA
• AMD
• NV
• QCOM
• WEBGPU

并且很容易添加更多！你选择的加速器只需支持大约 25 个底层操作即可。

要查看默认加速器，请运行：python3 -c "from tinygrad import Device; print(Device.DEFAULT)"

安装

目前推荐从源码安装 tinygrad。

从源码安装

git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

直接安装（master 分支）

python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

文档

详细的文档及快速入门指南可在 docs 网站 上找到，该网站基于 docs/ 目录构建而成。

快速对比 PyTorch 的示例

from tinygrad 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.tolist())  # dz/dx
print(y.grad.tolist())  # dz/dy

同样的代码在 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

贡献指南

近来，tinygrad 引起了广泛关注。遵循以下指南将有助于您的 PR 被接受。

首先，我们列出一些会导致 PR 直接关闭并指向本节的情况：

• 请勿进行代码极简化！虽然低行数是本项目的指导原则之一，但任何看起来像代码极简化的提交都会被直接关闭。我们的真正目标是降低复杂度、提高可读性，而单纯删除换行符并不能达到这一目的。
• 所有文档和空格变更的 PR 都会被关闭，除非您是广为人知的贡献者。编写文档的人应当是对代码库最为熟悉的开发者。未展现出相关能力的人员不应随意修改文档。此外，空格变更既无实际意义，还可能引入 bug。
• 您声称的“性能提升”必须经过基准测试验证。总体而言，本项目追求简洁性，因此即使您的 PR 使代码运行速度略有提升，也需权衡维护性和可读性的影响。
• 通常情况下，tinygrad/ 核心目录之外的代码未经充分测试，因此除非现有代码存在明显问题，否则不建议对其进行修改。
• 如果您的 PR 看起来过于复杂、改动量过大或新增了大量代码行，则不会被评审或合并。建议将其拆分为多个更小、更清晰的 PR。常见模式是在添加新功能之前先进行必要的重构。如果能够通过（干净地）重构，使新功能仅需三行代码即可实现，这将是理想的方案，也便于我们快速评审。

接下来，我们欢迎以下类型的贡献：

• Bug 修复（附带回归测试）非常受欢迎！目前该库尚未发布 1.0 版本，如果您发现了 bug，请修复它、编写相应的测试用例，并提交 PR，这将是非常有价值的贡献。
• 解决悬赏任务！tinygrad 为某些库的改进提供了现金悬赏（详见 Google 表格）。所有新增代码都应具备高质量并经过充分测试。
• 新功能。不过，在添加新功能时，请考虑代码行数与实用性之间的平衡。如果是三行代码的功能，其满足实用性的门槛自然低于三十或三百行的功能。所有新功能都必须包含回归测试。在没有其他约束的情况下，新功能的 API 应尽量与 PyTorch 或 NumPy 保持一致。
• 明显有益的重构。一般来说，若您的重构无法带来明确的优势，PR 将被关闭。然而，有些重构确实非常出色！请从深层次上考虑代码的可读性。仅仅调整空格或简单移动几个函数并无意义，但如果发现两个各有一百行的函数实际上可以使用同一个一百一十行的函数，并通过传递参数来实现，同时还能显著提升可读性，那么这就是一次巨大的改进。此类重构应通过 进程回放测试。
• 测试用例和模糊测试工具。如果您能添加不易出错的测试用例，我们将非常欢迎。我们目前也有一些模糊测试工具，借助它们以及进一步优化这些工具，仍有许多潜在的 bug 可以被发现。即使只是编写一些本应失败但预期会通过的测试用例（使用 @unittest.expectedFailure），也是非常有价值的。正是通过这种方式，我们才能不断取得进展。
• 移除核心 tinygrad/ 目录中的死代码。我们并不关心额外模块中的代码，但从核心库中移除死代码是非常有益的，这样新加入的开发者就不必阅读和困惑于这些无用的代码了。

运行测试

请使用 pre-commit install 安装预提交钩子。这将在每次提交时自动运行 linter、mypy 以及部分测试。

如需更多关于如何运行完整测试套件的示例，请参阅 CI 工作流。

以下是一些本地运行测试的示例：

python3 -m pip install -e '.[testing]'  # 安装测试所需的额外依赖
python3 test/backend/test_ops.py        # 仅运行算子测试
python3 -m pytest test/                 # 运行整个测试套件

进程回放测试

进程回放 会将您 PR 生成的内核与主分支的内核进行对比。如果您的 PR 是一次重构或性能优化，且未涉及行为上的变化，则应在 PR 标题中注明 [pr]。

tinygrad 快速上手指南

tinygrad 是一个端到端的深度学习栈，定位介于 PyTorch 和 micrograd 之间。它拥有自动微分张量库、融合内核的编译器、JIT 执行以及完整的神经网络训练组件，同时保持代码极简且易于黑客式修改。

环境准备

• 操作系统：Linux, macOS, Windows (WSL 推荐)
• Python 版本：Python 3.8 或更高版本
• 前置依赖：
  • git：用于克隆源码
  • pip：Python 包管理工具
  • (可选) 特定硬件驱动：如需使用 GPU 加速，请确保已安装对应的 CUDA、ROCm、Metal 或 OpenCL 驱动。

安装步骤

推荐从源码安装以获取最新功能，同时也支持直接通过 pip 安装主分支版本。

方式一：从源码安装（推荐）

git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .

方式二：直接安装（Master 分支）

python3 -m pip install git+https://github.com/tinygrad/tinygrad.git

提示：国内用户若遇到克隆或下载速度慢的问题，可配置 Git 和 Pip 的国内镜像源（如清华源、阿里源）进行加速。

验证默认加速器

安装完成后，可运行以下命令查看当前默认使用的计算设备：

python3 -c "from tinygrad import Device; print(Device.DEFAULT)"

基本使用

tinygrad 的 API 设计与 PyTorch 高度相似，支持动态图执行和自动微分。

1. 简单的自动微分示例

以下示例演示了如何创建张量、执行矩阵乘法、求和并进行反向传播：

from tinygrad 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.tolist())  # dz/dx
print(y.grad.tolist())  # dz/dy

2. 构建并训练一个简单的神经网络

tinygrad 内置了 nn 模块，包含优化器、层定义等工具。以下是一个简单的线性网络训练循环：

from tinygrad import Tensor, nn

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)

# 模拟输入数据 (batch_size=4)
x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])

with Tensor.train():
  for i in range(10):
    optim.zero_grad()
    loss = model(x).sparse_categorical_crossentropy(y).backward()
    optim.step()
    print(i, loss.item())

3. 体验惰性计算 (Laziness)

tinygrad 默认采用惰性执行策略，操作会被融合为单个内核。可以通过设置 DEBUG 环境变量观察生成的代码：

# 查看执行过程
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()"

# 查看生成的底层代码
DEBUG=4 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()"