数据截至 2026-08-27(上游 commit 118e812b316f)

01 · Tensor 与符号化 autograd

这一章讲什么: 用户看得见的那层——Tensor API 如何只建图不计算,以及 .backward() 时梯度图是怎么被「现算」出来的。读完你会明白:为什么 tinygrad 没有 PyTorch 那种 autograd.Function 子类体系。

---

1. 它要解决的小问题

一个深度学习框架的前端要同时满足两个互相矛盾的要求:

• 用起来是 eager 的:a @ b + c 一行行写,像 NumPy 一样;
• 跑起来是编译的:这一串算子得融合成尽量少的 GPU kernel,还得能对它做反向传播。

PyTorch 的解法是 eager 立即算 + 每个算子挂一个 C++ backward 节点;tinygrad 的解法是先建图、后计算——把「算子调用」和「计算」彻底分开,反向传播也在这张图上做。

---

2. Tensor:UOp 的薄壳

2.1 全部家当就三个字段

Tensor 的 __slots__ 只有三个字段(tinygrad/tensor.py:281):

__slots__ = "uop", "is_param", "grad"

字段	含义
uop	这个张量对应的计算图节点(一个 UOp,详见第 2 章)
grad	.backward() 之后填进来的梯度张量
is_param	是否「参数」(优化器只更新 is_param 为真的张量,tinygrad/nn/optim.py:13)

没有 data 指针、没有 shape 缓存、没有 device 字段——这些全是 self.uop 上的 property(tinygrad/tensor.py:368-375)。张量的「值」此刻可能根本不存在,存在的只是「怎么算出它」的图。

2.2 构造:万物归一为 UOp

Tensor.__init__(tinygrad/tensor.py:283-325)做的事是把各种输入统一变成 UOp:

输入	变成什么
None / Python 标量	UOp.const(...)——图里的常量节点
list / tuple / bytes	UOp._frompy(...)——先建 PYTHON 设备 buffer
numpy 数组	_fromnp 建 NPY buffer 并「假装已 realize」(tinygrad/tensor.py:264-268)
pathlib.Path	DISK:<路径> 设备的 buffer(tinygrad/tensor.py:286、:318)——权重文件直接就是张量
UOp	直接用,按需 copy_to_device / cast(tinygrad/tensor.py:320-322)

构造成功后,每个 Tensor 都会被登记进全局弱引用表 all_tensors(tinygrad/tensor.py:325)。这张表后面有两个用处:realize 后把新 UOp 回填给所有存活张量(_apply_map_to_tensors,tinygrad/tensor.py:244),以及 backward 时找出「需要梯度的全体张量」(tinygrad/tensor.py:665-666)。

---

3. 算子即建图

3.1 mixin 体系

Tensor 的几百个算子不在 tensor.py 里,而是按类别拆在 tinygrad/mixin/ 下。单链继承,每层加一类能力(tinygrad/mixin/):

MovementMixin ─► CreationMixin ─► ElementwiseMixin ─► OpMixin ─► RandMixin ─► UOp / Tensor
 移动/视图         创建/ dtype       逐元素算术         + ReduceMixin   随机数
 (reshape 等)     (empty/full)     (加减乘除)         (sum/max/…)   (rand/randn)

• ElementwiseMixin(tinygrad/mixin/elementwise.py:12):加减乘除、比较、激活。核心抽象只有一个 alu(op, *src),其余全是它的语法糖,例如 _binop(tinygrad/mixin/elementwise.py:33-36)先做广播和 dtype 提升,再调 alu。
• MovementMixin:reshape/permute/expand/pad/shrink/flip,统一走 _mop(tinygrad/tensor.py:677)。
• ReduceMixin:sum/max/prod,统一走 _rop(tinygrad/tensor.py:678)。
• OpMixin(tinygrad/mixin/op.py:19):matmul、conv、einsum 等复合算子,全部用上面三类拼出来。

UOp 自己也继承这条链(class UOp(RandMixin, ...),tinygrad/uop/ops.py:238),所以重写规则里拼新子图和写 Tensor 代码是同一套 API——这是 tinygrad 代码密度高的一个原因。

3.2 一次调用的全部动作:_apply_uop

所有算子最终汇到一个十行函数(tinygrad/tensor.py:330-339):

def _apply_uop(self, fxn, *x, **kwargs):
  srcs = (self,)+x
  new_uop = fxn(*[t.uop for t in srcs], **kwargs)   # 在各个输入的 uop 上建一个新节点
  ...
  ret = Tensor.__new__(Tensor)                       # 绕过 __init__,直接包回去
  ret.uop, ret.grad, ret.is_param = new_uop, None, True
  all_tensors[weakref.ref(ret)] = None
  return ret

重点看:没有任何计算。a + b 只是造出一个 UOp(Ops.ADD, src=(a.uop, b.uop))。

3.3 两类算子,两种命运

UOp 图里的算子可以粗分两群(这个划分是第 3 章调度器的全部前提):

群	成员	命运
movement(移动)	RESHAPE/PERMUTE/EXPAND/PAD/SHRINK/FLIP(GroupOp.Movement,tinygrad/uop/__init__.py:120)	永远不落盘,只改「坐标怎么读」
compute(计算)	ADD/MUL/...(GroupOp.ALU)、REDUCE、WMMA、LOAD/STORE	真正产生数据,要进 kernel

Tensor.dot(tinygrad/mixin/op.py:364-387)是最好的例子:一个 matmul 就是 reshape + 转置 + 广播乘 + sum 四种基本节点的组合,前端没有专门的 "MatMul 算子"。

---

4. 符号化 autograd

4.1 它要解决的小问题

反向传播需要给每个算子定义「梯度怎么往回传」。传统做法(PyTorch/micrograd)是给每个算子配一个 backward 函数,前向时记住它;tinygrad 的做法是:不存任何 backward 闭包,需要时对 UOp 图现场套规则。因为图就在那儿,规则也在那儿,梯度图现算就行。

4.2 直觉:梯度规则就是一张模式匹配表

pm_gradient(tinygrad/mixin/gradient.py:55-91)是一张 UPat → 梯度子图 的表,每个规则拿到「回传来的梯度 ctx」和「前向节点 ret」,返回对各个输入的梯度:

pm_gradient = PatternMatcher([
  (UPat(Ops.ADD), lambda ctx: (ctx, ctx)),                                   # d(a+b): 各回传一份
  (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),  # d(a*b)
  (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
  (UPat(Ops.REDUCE, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg[0])),
  (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
  (UPat((Ops.CMPLT, Ops.CMPNE)), lambda: (None, None)),                      # 比较不可导
  ...
])

注意它不是 Python 回调,而就是重写规则——返回的梯度是用 UOp 拼出来的新子图。

4.3 图示:一次 backward 的三步

 z = (x*w).sum()            前向图(UOp)
      │ z.backward()
      ▼
 ① compute_gradient:从 z 倒序走图,逐节点套 pm_gradient
      │   得到的也是 UOp 图:dz/dx、dz/dw
      ▼
 ② 广播边自动求和:分支点把多份梯度 ADD 起来
      │
      ▼
 ③ 梯度 Tensor 写回 x.grad / w.grad
    (之后随下一次 realize 一起被调度、融合、编译)

4.4 原理演示

# 示意,非源码 —— 演示「套规则」的逆序遍历
grads = {z: ones_like(z)}            # 根的梯度是 1
for node in reversed(toposort(z)):   # 从输出往输入走
    g = grads.get(node)
    if g is None: continue
    for src, local in zip(node.src, grad_rules[node.op](g, node)):
        if local is not None:
            grads[src] = grads.get(src, 0) + local   # 分支点:梯度相加

重点看最后那行梯度相加:这是 autograd 的数学核心(多元链式法则),micrograd 里它藏在 __add__ 里,这里它是一个显式的 +(tinygrad/mixin/gradient.py:132)。

4.5 真实实现

compute_gradient(tinygrad/mixin/gradient.py:100-139)就是上面伪代码的正式版,多了三件事:

1. 只走需要的子图:_deepwalk(tinygrad/mixin/gradient.py:93-98)先算出「在目标路径上」的节点,DETACH 直接截断。
2. 广播形状对齐:边的两端形状不一致时,沿广播轴把梯度 sum 回源形状(tinygrad/mixin/gradient.py:126-127,用 broadcast_axes 找出被广播掉的轴)。
3. 跨 CALL/FUNCTION 边界:遇到 @function 产生的节点,调 call_gradient(tinygrad/mixin/gradient.py:23-52)——把反向也打包成一个新的 _backward CALL,支持 precompile_backward。

入口两个,殊途同归:

• Tensor.gradient(*targets)(tinygrad/mixin/op.py:450-471):要求浮点、默认根梯度 1.0,返回各目标的梯度;
• Tensor.backward()(tinygrad/tensor.py:653-673):在 gradient 之上,从 all_tensors 里圈出「本图涉及的浮点张量」,把梯度写进 .grad,已有梯度则 assign 累加。

4.6 这带来一个免费红利

因为反向图也是 UOp 图,它和前向图走完全相同的 rangeify 融合、BEAM 优化、编译管线。PyTorch 里「反向算子」是另一套预先写好的 CUDA kernel;tinygrad 里反向 kernel 是编译器现场生成的,和前向一样可以融合。README 说的「IR-based autodiff over primitives (like JAXPR + XLA)」指的就是这个。

---

5. @function:把一段 Python 变成一个图节点

@TinyJit 管「重放」,@function(tinygrad/function.py:32-105)管「抽象」:把任意 Python 函数的正向变成一个可参数化的 FUNCTION/CALL 节点。

_function.__call__(tinygrad/function.py:43-90)的流程:

1. 收集所有入参张量的 uop,去重(tinygrad/function.py:46-49);
2. 禁用设备,原样执行函数体(tinygrad/function.py:53-58)——此刻只建图;
3. 把已知输入替换成 PARAM 槽位(tinygrad/function.py:67-68);
4. 图里剩下的 BUFFER 就是隐式输入(函数闭包捕获的权重),pm_ctx(tinygrad/function.py:15-19)把它们也参数化;allow_implicit=False 时会显式报出来(tinygrad/function.py:75-79);
5. 返回 uret.call(*call_uops, ...)——之后整段前向在图里就是一个节点,可以被调度器整体处理,也可以像 examples/beautiful_mnist.py:18 那样给整个模型 __call__ 加上。

precompile=True 时,正向被提前编译好,反向通过 grad_fxn 或自动求导挂接(call_gradient 里 k.arg.precompile 分支,tinygrad/mixin/gradient.py:43-44)。

---

6. 关键细节与坑

• weak dtype 没有存储。 Python 标量造出的常量是 weakint/weakfloat(dtype 推导见 tinygrad/uop/ops.py:181-187),参与运算时自动提升;但你不能 realize 一个 weak dtype 张量(tinygrad/tensor.py:403-404),也不能对它 __setitem__(tinygrad/tensor.py:681)。
• assign 是图里的副作用。 原地写不是改数据,而是给 uop 包一层 AFTER(STORE)(tinygrad/tensor.py:451);依赖这层 AFTER,调度器才知道谁必须先算(第 3 章的 RAW/WAR 全靠它)。
• __setitem__ 是前端最拧巴的地方。 它要分辨「纯视图的原位写」「高级索引」「惰图上的写」,还得拒绝「别的张量还引用着同一块底层 buffer」的写法(tinygrad/tensor.py:685-691 的 "can't setitem on a tensor with other uses")。这段逻辑基本全是坑堆出来的,读它之前先读测试。
• Tensor.__bool__ 直接抛错(tinygrad/tensor.py:362):张量没有真值,if t: 是 bug 的信号,不是分支。
• 梯度只对浮点存在。 gradient 里显式检查(tinygrad/mixin/op.py:465),整数张量不在 autograd 世界内。
• @function 里禁止 realize。 函数体在建图期跑,ALLOW_DEVICE_USAGE 被关掉(tinygrad/function.py:53),在函数里调用会碰设备数据的 API(如 .numpy())会炸。

---

7. 代码地图

主题	文件路径	符号名
Tensor 定义与构造	tinygrad/tensor.py	Tensor.__slots__、Tensor.__init__、all_tensors
建图枢纽	tinygrad/tensor.py	Tensor._apply_uop、Tensor.alu、Tensor._mop、Tensor._rop
算子库	tinygrad/mixin/elementwise.py、tinygrad/mixin/movement.py、tinygrad/mixin/op.py	ElementwiseMixin.alu、MovementMixin、OpMixin.dot
autograd 规则表	tinygrad/mixin/gradient.py	pm_gradient、reduce_gradient
autograd 遍历	tinygrad/mixin/gradient.py	compute_gradient、_deepwalk、call_gradient
backward 入口	tinygrad/tensor.py、tinygrad/mixin/op.py	Tensor.backward、OpMixin.gradient
函数抽象	tinygrad/function.py	function、_function.__call__、pm_ctx
优化器衔接	tinygrad/nn/optim.py	Optimizer.schedule_step、fused optimizer