tinygrad Tensor 高阶算子完全指南：Reduce、卷积、注意力与线性代数 API 深度解析

tinygrad Tensor 高阶算子完全指南：Reduce、卷积、注意力与线性代数 API 深度解析

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad

本文以 tinygrad 官方 API 文档 docs/tensor/ops.md 为核心骨架，系统讲解Tensor在规约（Reduce）、处理（Processing）、神经网络函数式算子（Neural Network functional）与线性代数（Linear Algebra）四大类中的高阶算子。每个算子均结合源码实现、参数语义与可执行示例展开，读者读完可以准确掌握这些算子的行为差异、数值稳定性处理与底层实现原理，并能直接用 tinygrad 编写卷积网络、注意力机制与损失函数。

docs/tensor/ops.md是 tinygrad 文档站 "Complex Ops"（复杂算子）章节的源文件，它本身不写代码，而是通过 mkdocstrings 插件以::: tinygrad.Tensor.xxx指令从源码 docstring 自动渲染出完整的 API 文档（见 mkdocs.yml 中mkdocstrings的配置）。因此本文的内容主体实质上是Tensor各算子 docstring 与实现本身。这些算子按用途分布在源码中：规约类基础算子位于 tinygrad/mixin/reduce.py，其余大部分位于 tinygrad/mixin/op.py，随机/注意力类位于 tinygrad/mixin/rand.py，逐元素比较类位于 tinygrad/mixin/elementwise.py，最终统一挂载到 tinygrad/tensor.py 的Tensor类上。

一、Reduce：规约与归一化算子

规约算子把张量沿一个或多个轴压缩为统计量。tinygrad 的统一入口是_reduce(op, axis, keepdim)，它会先把axis归一化为元组，再调用底层_rop完成规约（见 tinygrad/mixin/reduce.py）：

def _reduce(self, op:Ops, axis:int|Sequence[int]|None=None, keepdim=False) -> Self: self = self.cast(self.commit_dtype()) axis = tuple(self._resolve_dim(x) for x in (range(self.ndim) if axis is None else make_tuple(axis, 1))) if self.ndim == 0: axis = () ret = self._rop(op, axis) return ret.reshape(tuple(1 if i in axis else s for i,s in enumerate(self.shape))) if keepdim else ret

• axis：可传单个整数、整数序列或None。axis=None表示对所有维度规约；keepdim=True时被规约的维度保留为长度 1。
• 实现细节：axis支持负数（_resolve_dim会解析为实际维号）；标量（0 维）张量的axis会被强制置空。

1. sum / prod / max / min / any / all

这六个算子共享axis与keepdim两个关键词参数，其中sum与prod还额外支持dtype参数控制累加的数据类型：

• sum(axis=None, keepdim=False, dtype=None)：沿指定轴求和。源码中当dtype未指定时，会先 cast 到sum_acc_dtype(commit_dtype())（按输入 dtype 选择的累加类型，例如 float16/bfloat16/fp8 会提升精度），累加完成后再还原为原 dtype（tinygrad/mixin/reduce.py）。这一设计避免了低精度累加带来的数值误差。
• prod(axis=None, keepdim=False, dtype=None)：沿指定轴求积，默认按输入自身 dtype 累乘。
• max(axis=None, keepdim=False)：沿轴取最大值；min在 tinygrad/mixin/op.py 中通过_inverse().max(...)._inverse()（取负后求最大再取负）复用max实现。
• any(axis=None, keepdim=False)：沿轴测试"是否存在 True"，实现为self.bool().max(axis, keepdim)——布尔化后取最大值，等价于逻辑或。
• all(axis=None, keepdim=False)：沿轴测试"是否全部为 True"，实现为self.bool().prod(axis, keepdim)——布尔化后求积，等价于逻辑与。

t = Tensor.arange(6).reshape(2, 3) print(t.sum().numpy()) # 15 print(t.sum(axis=0).numpy()) # [3, 5, 7] print(t.sum(axis=1).numpy()) # [3, 12] print(t.prod(axis=0).numpy()) # [0, 4, 10] print(Tensor([[1, 0, 2], [5, 4, 3]]).max().numpy()) # 5 print(Tensor([[1, 0, 2], [5, 4, 3]]).max(axis=1, keepdim=True).numpy()) # [[2], [5]]

2. mean / var / std 与组合算子

• mean(axis=None, keepdim=False)：算术平均。源码（tinygrad/mixin/op.py）先以累加精度求和，再除以被规约元素个数prod(...)，最后 cast 回输出 dtype（整数输入输出 float32）。
• var(axis=None, keepdim=False, correction=1)：方差。实现为(self - mean).square()求和后除以n - correction，其中correction=1对应 Bessel 校正（样本方差，除数smax(n-1, 0)保证不为负）；correction=0则为总体方差。
• std(axis=None, keepdim=False, correction=1)：标准差，即var(...).sqrt()。
• var_mean/std_mean：一次计算同时返回(方差, 均值)/(标准差, 均值)二元组，与 PyTorch 的torch.var_mean/torch.std_mean对齐，是各自的语法糖封装（tinygrad/mixin/op.py）。

Tensor.manual_seed(42) t = Tensor.normal(2, 3, mean=2.5, std=0.5) print(t.var().numpy()) # 总体/样本方差（默认 correction=1） print(t.var(axis=0).numpy()) # 沿轴 0 的方差 var, mean = t.var_mean() std, mean = t.std_mean()

3. softmax / log_softmax / logsumexp / logcumsumexp

这组算子核心是数值稳定性处理。_softmax的实现（tinygrad/mixin/op.py）先减去沿轴最大值（self.max(axis, keepdim=True).detach()）再做指数运算，避免大数值溢出：

def _softmax(self, axis, dtype=None) -> tuple[Self, Self, Self]: m = self - self.max(axis=axis, keepdim=True).detach() if dtype is not None: m = m.cast(to_dtype(dtype)) e = m.exp() return m, e, e.sum(axis=axis, keepdim=True)

• softmax(axis=-1, dtype=None)：输出落在[0, 1]且沿轴求和为 1，公式为e * sum.reciprocal()（用倒数而非除法）。
• log_softmax(axis=-1, dtype=None)：m - ss.log()，即对 softmax 取对数，同样稳定。
• logsumexp(axis=None, keepdim=False)：log(sum(exp(x)))的稳定形式。源码先取mx = max(axis, keepdim=True).detach()，把不可靠的最大值替换为 0（isfinite().where(mx, 0)），再计算(self - m).exp().sum().log() + m（tinygrad/mixin/op.py）。注意detach()意味着 logsumexp 的梯度只流经减最大值之后的路径。
• logcumsumexp(axis=0)：log(cumsum(exp(x)))的稳定实现，使用cummax代替逐个最大值来稳定指数累加（tinygrad/mixin/op.py），0 维输入直接原样返回。

Tensor.manual_seed(42) t = Tensor.randn(2, 3) print(t.softmax().numpy()) # 默认沿最后一维 print(t.softmax(axis=0).numpy()) print(t.log_softmax(axis=1).numpy()) print(t.logsumexp(axis=1).numpy()) print(t.logcumsumexp(axis=1).numpy())

4. argmax / argmin 与 isclose / allclose

• argmax(axis=None, keepdim=False)：返回最大值下标。axis=None时先flatten()再求全局下标；源码通过"最大值掩码 × 逆序 arange → 取 max → 用轴长度相减"的技巧把"第一个最大值"的下标还原为正序（tinygrad/mixin/op.py），输出 dtype 为 int32。argmin复用_inverse().argmax(...)（tinygrad/mixin/op.py）。
• isclose(other, rtol=1e-05, atol=1e-08, equal_nan=False)：逐元素容差比较，位于 tinygrad/mixin/elementwise.py。判定分三部分：有限值满足|self-other| <= atol + rtol*|other|；无穷值必须相等；equal_nan=True时 NaN 视为相等。
• allclose(other, rtol=1e-05, atol=1e-08, equal_nan=False)：isclose(...).all()的包装，返回单个标量布尔张量（tinygrad/mixin/op.py），常用于数值正确性断言。

t = Tensor([[1, 0, 2], [5, 4, 3]]) print(t.argmax().numpy()) # 3（展平后最大值的下标） print(t.argmax(axis=0).numpy()) # [1, 1, 1] print(t.argmax(axis=1).numpy()) # [2, 0] print(Tensor([1e-7, float('nan')]).isclose(Tensor([0.0, float('nan')]), equal_nan=True).numpy()) # [True, True]

二、Processing：卷积、矩阵乘法与张量处理算子

1. conv2d / conv_transpose2d 与池化：三种 padding 形式

conv2d、conv_transpose2d、avg_pool2d、max_pool2d、max_unpool2d这组算子统一支持三种 padding 写法（以padding参数为例）：

1. 单个int：所有空间维度统一使用该填充值；
2. 长度为空间维度数的元组：(padding_height, padding_width, ...)，每个空间维一个值；
3. 长度为空间维度数两倍的元组：(padding_left, padding_right, padding_top, padding_bottom, ...)，逐边指定。

• conv2d(weight, bias=None, groups=1, stride=1, dilation=1, padding=0, dtype=None)：对self施加卷积。输入形状(bs, cin, ...)，权重形状(cout, cin, *kernel)，要求groups*cin == 输入通道数。源码（tinygrad/mixin/op.py）先把输入_pad_constant补零，再用_pool做滑窗展开，最后与权重广播相乘并沿(cin, *kernel)求和得到输出——卷积在 tinygrad 中被表达为一次带 padding 的池化加一次规约。当IMAGE环境启用时走image_conv2d（图像内存格式优化路径），当WINO启用且核为 3×3、stride=dilation=1 时走_conv2d_winograd（Winograd F(4x4,3x3) 变换，见 tinygrad/mixin/op.py）。
• conv_transpose2d(weight, bias=None, groups=1, stride=1, dilation=1, padding=0, output_padding=0)：转置卷积。实现先把 stride 展开为"插入零行"（reshape→pad→reshape→shrink），再反算等效 padding 后复用conv2d（tinygrad/mixin/op.py）。
• avg_pool2d(kernel_size=(2,2), stride=None, dilation=1, padding=0, ceil_mode=False, count_include_pad=True)：平均池化。count_include_pad=False时把 padding 区域排除在分母之外（源码用"全 1 张量同样池化求计数"的方式实现，tinygrad/mixin/op.py）；ceil_mode=True时输出形状按向上取整计算。
• max_pool2d(kernel_size=(2,2), stride=None, dilation=1, padding=0, ceil_mode=False, return_indices=False)：最大池化，padding 区填dtype.min后取最大；return_indices=True时同时返回 argmax 下标，供max_unpool2d使用。
• max_unpool2d(indices, kernel_size=(2,2), ...)：max_pool2d的部分逆运算——把池化输出放回indices指定的原始位置，其余位置填 0；output_size可显式指定输出形状消除歧义（tinygrad/mixin/op.py）。

t = Tensor.arange(25).reshape(1, 1, 5, 5) print(t.avg_pool2d().numpy()) # 默认 2x2 print(t.avg_pool2d(padding=1, count_include_pad=False).numpy()) print(t.max_pool2d(ceil_mode=True).numpy()) output, indices = Tensor.max_pool2d(t, return_indices=True) print(Tensor.max_unpool2d(output, indices).numpy()) # 卷积：输入 (1,1,3,3)，权重 (1,1,2,2) t = Tensor.arange(9).reshape(1, 1, 3, 3) w = Tensor.ones(1, 1, 2, 2) print(t.conv2d(w).numpy())

与 PyTorch 的一个显著差异：以上池化与卷积实现不局限于 2D，传入任意数量的空间维度即可推广到 3D、4D 等（docstring 中明确标注 "unlike PyTorch, this implementation is not limited to only 2d"）。

2. dot / matmul / einsum

• dot(w, dtype=None)：点积。w为 1 维时是self最后轴与w的求和积；w为 N 维时是self最后轴与w倒数第二轴的求和积。实现通过 reshape/transpose 对齐后执行(x*w).sum(-1, dtype=dtype)，最后 cast 回两者的最小上界 dtype（tinygrad/mixin/op.py）。
• matmul(x, reverse=False, dtype=None)：矩阵乘法，等价于self.dot(x)；reverse=True时执行x.dot(self)。运算符@（__matmul__）与右乘__rmatmul__均已重载。
• einsum(formula, *operands, dtype=None)：爱因斯坦求和约定，作为类方法调用，如Tensor.einsum("ij,ij->", x, y)返回逐元素乘积之和。实现（tinygrad/mixin/op.py）支持省略号...（自动展开为字母）、单输入内重复字母的迹（trace）、以及输出顺序重排：流程为"消重（取对角线）→ 对齐字母表 → 逐项相乘 → 对非输出字母求和 → 按输出顺序 permute"。

a = Tensor([[1, 2], [3, 4]]) b = Tensor([[5, 6], [7, 8]]) print(a.dot(b).numpy()) # [[19, 22], [43, 50]] print(a.matmul(b).numpy()) # 同上 print(a @ b) # 运算符形式 print(Tensor.einsum("ij,ij->", a, b).numpy()) # 70 print(Tensor.einsum("ij,jk->ik", a, b).numpy()) # 标准矩阵乘

3. 累积算子：cumsum / cumprod / cummax / cummin

• cumsum(axis=0)/cumprod(axis)：沿轴累积求和/求积。
• cummax(axis=0)/cummin(axis=0)：返回(values, indices)二元组，indices 是该位置累积最值的下标，输出为 int32。
• 底层实现（tinygrad/mixin/op.py）：通过_cumalu把累积运算转化为"pad 后_pool滑窗 + 对应规约"；当轴长度超过阈值SPLIT = 256时，_split_cumalu采用两阶段分块累积（先按 256 块内累积，再对块尾累积并广播叠加），避免长序列上的串行依赖。

t = Tensor.ones(2, 3) print(t.cumsum(1).numpy()) # [[1, 2, 3], [1, 2, 3]] t = Tensor([0, 1, -1, 2, -2, 3, -3]) values, indices = t.cummax(0) print(values.numpy()) # [0, 1, 1, 2, 2, 3, 3] print(indices.numpy()) # [0, 1, 1, 3, 3, 5, 5]

4. triu / tril、interpolate 与排序族

• triu(diagonal=0)/tril(diagonal=0)：取上/下三角，其余元素置 0。diagonal=0为主对角线；正数偏向上方、负数偏向下方。实现基于_tri掩码配合where选择（tinygrad/mixin/op.py）。
• interpolate(size, mode="linear", align_corners=False)：上/下采样到目标size，接受 0 到 N 个 batch 维。mode支持linear、nearest、nearest-exact；传 2D 或 3D 的size即实现 bilinear / trilinear。linear模式通过gather取低/高邻居再按小数部分lerp插值；align_corners=False（默认，与 PyTorch 一致的"像素对齐"语义）与True（角点对齐）使用不同的坐标映射公式（tinygrad/mixin/op.py）。align_corners仅对linear模式允许设置。
• sort(dim=-1, descending=False)：沿指定维做双调排序（bitonic sort），返回(sorted_values, indices)二元组；等值元素的顺序保持稳定。实现先把长度补齐到 2 的幂，再按双调排序器的 stage/substage 用split/cat/maximum/minimum/flip逐层归并（tinygrad/mixin/op.py）。
• argsort(dim=-1, descending=False)：仅返回排序下标，即sort(...)[1]。
• topk(k, dim=-1, largest=True, sorted_=True)：返回沿dim的最大（largest=True）/最小k个值及下标。当前不支持sorted_=False（直接抛NotImplementedError），实现为sort后shrink_to截取前 k 个（tinygrad/mixin/op.py）。

t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) print(t.triu(diagonal=0).numpy()) print(t.tril(diagonal=-1).numpy()) t = Tensor([[0.1, 0.5, 1.2, 3.4, 2.1], [2.2, 1.9, 0.3, 4.5, 0.8]]) sorted_values, indices = t.sort(dim=1, descending=True) print(sorted_values.numpy()) print(t.topk(2, dim=1).numpy()) # (values, indices) 二元组 print(t.argsort(dim=1).numpy())

5. scatter / scatter_reduce / masked_select / nonzero / multinomial

• scatter(dim, index, src, reduce=None)：把src按index写入self沿dim的位置。reduce支持None（覆盖写，重复下标取最后写入）、'add'、'multiply'。注意：带 reduce 且src为张量时需使用scatter_reduce；src为标量常量时可用 reduce（tinygrad/mixin/op.py）。
• scatter_reduce(dim, index, src, reduce, include_self=True)：reduce支持"sum" | "prod" | "mean" | "amax" | "amin"五种规约；include_self=False时把self的对应值排除在规约之外。实现通过_pre_scatter把index转成 one-hot 掩码，再在dim上做对应规约（tinygrad/mixin/op.py）。
• masked_select(mask, size=None, fill_value=0)：按布尔掩码取元素（展平后输出）。size=None时输出长度等于 True 个数，此时不可 JIT；传入size=N则输出定长N，不足用fill_value填充、超出截断，可 JIT（tinygrad/mixin/op.py）。要求mask为布尔 dtype。
• nonzero(size=None, fill_value=0)：返回非零元素下标，输出形状(n_nonzero, ndim)；同样支持size定长填充以启用 JIT（tinygrad/mixin/op.py）。
• multinomial(num_samples=1, replacement=False)：按self作为权重做多项式采样。replacement=True或num_samples==1时用"累计分布函数 + 均匀采样"（(unif_samples >= cdf).sum()）；replacement=False时采用Efraimidis-Spirakis 算法：(weight.rand_like().log2() / weight).topk(num_samples, dim=1)（tinygrad/mixin/rand.py）。self须为 1 维或 2 维，输出 int32。

src = Tensor.arange(1, 11).reshape(2, 5) index = Tensor([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]) print(Tensor.ones(1, 5).scatter_reduce(0, index, src, reduce='sum').numpy()) print(Tensor.zeros(3, 5).scatter(0, Tensor([[0, 1, 2, 0]]), src).numpy()) t = Tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) mask = Tensor([[True, False, True], [False, True, False], [False, False, True]]) print(t.masked_select(mask).numpy()) # [0, 2, 4, 8] print(t.masked_select(mask, size=6, fill_value=-1).numpy()) Tensor.manual_seed(42) print(Tensor([1, 2, 3, 4]).multinomial(20, replacement=True).numpy())

三、Neural Network (functional)：神经网络函数式算子

1. linear / sequential / layernorm / batchnorm / dropout / one_hot

• linear(weight, bias=None, dtype=None)：线性变换x @ W^T + b。权重为 1 维时退化为x * weight，否则self.dot(weight)，最后可选加 bias；传dtype时先统一 cast（tinygrad/mixin/op.py）。
• sequential(ll)：把一列Tensor -> Tensor函数依次串联，等价于functools.reduce(lambda x, f: f(x), ll, self)。
• layernorm(axis=-1, eps=1e-5)：Layer Normalization（tinygrad/mixin/op.py）。实现为(x - mean) * rsqrt(mean((x-mean)^2) + eps)，不含可学习的 weight/bias（与 PyTorch 的 elementwise_affine=False 语义对应）。
• batchnorm(weight, bias, mean, invstd, axis=1)：Batch Normalization 的推理形式——直接接收训练阶段统计出的mean与invstd（1/sqrt(var+eps)），weight/bias可为None（tinygrad/mixin/op.py）。典型用法见 docstring：t.batchnorm(None, None, t.mean(axis=(0,2,3)), t.var(axis=(0,2,3)).add(1e-5).rsqrt())。
• dropout(p=0.5)：仅在TRAINING为真时生效（Context(TRAINING=1)内），否则恒等返回。实现为(rand_like() >= p).where(self, 0) / (1-p)——经典 inverted dropout 缩放（tinygrad/mixin/rand.py）；p=0恒等，p=1全零，p越界抛ValueError。
• one_hot(num_classes)：把整数下标张量转 one-hot（输出比输入多一维，num_classes < 0或非整数 dtype 会报错）。实现为self[..., None]._one_hot_along_dim(num_classes).where(1, 0)（tinygrad/mixin/op.py）。

t = Tensor([[1, 2], [3, 4]]) weight = Tensor([[1, 2], [3, 4]]) bias = Tensor([1, 2]) print(t.linear(weight, bias).numpy()) print(t.sequential([lambda x: x * 2, lambda x: x + 1]).numpy()) t = Tensor.randn(8, 10, 16) * 2 + 8 t = t.layernorm() print(t.mean().item(), t.std().item()) # 约 0, 1 Tensor.manual_seed(42) t = Tensor.randn(2, 2) with Context(TRAINING=1): print(t.dropout().numpy()) print(Tensor([0, 1, 3, 3, 4]).one_hot(5).numpy())

2. scaled_dot_product_attention（SDPA）

scaled_dot_product_attention(key, value, attn_mask=None, dropout_p=0.0, is_causal=False, enable_gqa=False)是 tinygrad/mixin/rand.py 中的标准缩放点积注意力实现，self为 query：

qk = q.matmul(key.transpose(-2,-1), dtype=least_upper_dtype(q.dtype, key.dtype, dtypes.float32)) / math.sqrt(q.shape[-1]) if is_causal: if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True") attn_mask = qk.const_like(True, dtypes.bool).tril() if attn_mask is not None: if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf")) qk = qk + attn_mask return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value

要点：QK 点积强制以 float32 及以上精度计算（least_upper_dtype(q, k, float32)）并除以sqrt(head_dim)；is_causal=True时自动生成下三角掩码，且与attn_mask互斥；布尔掩码会被转换为 0/-inf 加法掩码（而非乘法掩码）；dropout_p经由dropout实现（同样仅在 TRAINING 下生效）；enable_gqa=True时通过repeat_interleave把 key/value 的头数扩展到与 query 一致，从而支持 Grouped Query Attention。

q = Tensor.randn(2, 4, 8) k = Tensor.randn(2, 4, 8) v = Tensor.randn(2, 4, 8) print(q.scaled_dot_product_attention(k, v).numpy()) print(q.scaled_dot_product_attention(k, v, is_causal=True).numpy())

3. 损失函数族

tinygrad 提供五个函数式损失，均支持reduction参数（"mean"默认 /"sum"/"none"，由_do_reduction统一处理，tinygrad/mixin/op.py）：

• binary_crossentropy(Y, reduction="mean")：二分类交叉熵-Y*log(p) - (1-Y)*log(1-p)，self为概率（对应torch.nn.BCELoss）。
• binary_crossentropy_logits(Y, reduction="mean", pos_weight=None)：self为 logits 版本，用logsigmoid保证数值稳定（对应BCEWithLogitsLoss）；pos_weight可调整正类权重。
• sparse_categorical_crossentropy(Y, ignore_index=-1, label_smoothing=0.0, reduction="mean")：self为 logits、Y为类别标签的交叉熵（对应CrossEntropyLoss的 sparse 形式）。注意与 PyTorch 不同：类别轴必须是最后一维（-1）；支持ignore_index屏蔽与label_smoothing平滑。
• cross_entropy(Y, reduction="mean", label_smoothing=0.0)：self为 logits，Y可为类别标签或类别概率分布；当Y.shape != self.shape时自动 one-hot 化。类别轴取 0（1 维输入）或 1（2 维输入）。
• nll_loss(Y, weight=None, ignore_index=None, reduction="mean")：负对数似然，self为 log-probabilities（配合log_softmax使用），Y为标签；weight可对类别加权。

t = Tensor([[-1, 2, -3], [1, -2, 3]]) Y = Tensor([1, 2]) print(t.cross_entropy(Y).item()) # 0.313... print(t.cross_entropy(Y, reduction='none').numpy()) print(t.sparse_categorical_crossentropy(Y).item()) print(t.log_softmax().nll_loss(Y).item()) # 与 cross_entropy 数值一致 t = Tensor([0.1, 0.9, 0.2]); Y = Tensor([0, 1, 0]) print(t.binary_crossentropy(Y).item()) t = Tensor([-1, 2, -3]); Y = Tensor([0, 1, 0]) print(t.binary_crossentropy_logits(Y).item())

四、Linear Algebra：qr 与 svd

• qr()：QR 分解，返回(Q, R)，要求ndim > 1。实现采用逐列 Householder 反射：每轮构造长度为 m 的反射向量v与 rank-1 更新因子w = tau*v，迭代min(m, n)次更新R -= w @ (v.T @ R)、Q -= (Q @ v) @ w.T，支持 batch 维度（tinygrad/mixin/op.py）。
• svd(full_matrices=True)：奇异值分解，返回(U, S, Vh)，要求ndim > 1。实现基于经典 Jacobi 旋转法（源码注释引用 LAPACK lawn169 论文第 26 页）：先 QR 预处理，再以"轮转配对（round robin pairing）+ 2×2 Jacobi 旋转"迭代4*num轮收敛，最后对奇异值排序并重组 U/V（tinygrad/mixin/op.py）。full_matrices=False时返回瘦 SVD（U 只保留前num = min(m,n)列）。

t = Tensor.randn(3, 3) Q, R = t.qr() U, S, Vh = t.svd() print(U.shape, S.shape, Vh.shape)

五、测试验证：这些算子的行为边界

tinygrad 在 test/backend/test_ops.py 中为上述算子提供了系统性的数值验证，是理解行为边界的最佳参考：

• 规约类：test_sum、test_sum_dtype_arg（验证sum的 dtype 累加参数）、test_sum_collapse（验证规约与 view 折叠优化）、test_softmax、test_softmax_other_axis、test_softmax_argmax；
• 矩阵运算：test_matmul、test_matmul_batched、test_einsum、test_einsum_ellipsis、test_einsum_trace、test_einsum_shape_check、test_einsum_arity_check1/2（验证 operand 数量与 shape 的校验逻辑）；
• 卷积与池化：test_conv2d、test_conv2d_bs_4_cin_3、test_conv2d_errors（验证groups*cin约束）；
• 散射与注意力：test_scatter、test_scatter_add、test_scatter_mul、test_scatter_reduce、test_scatter_reduce_prod_zeros、test_scatter_reduce_errors、test_scaled_dot_product_attention、test_topk。

例如test_einsum_arity_check系列印证了 tinygrad/mixin/op.py 中len(xs) != len(inputs)会抛出ValueError的约束；test_conv2d_errors对应groups*cin == cin_的断言。运行方式为pytest test/backend/test_ops.py -k "sum or matmul or softmax"。

六、小结：tinygrad 高阶算子的设计哲学

纵观 docs/tensor/ops.md 覆盖的全部算子，可以提炼出 tinygrad 的几个鲜明设计特征：

1. 算子由少量基础 UOp 组合而成：卷积 = padding +_pool+ 规约；累积运算 = padding +_pool+ 对应 reduce；min=_inverse().max()._inverse()。这让后端代码生成器只需处理极少数底层算子，就能支持全部高阶 API。
2. 数值稳定性内建：softmax / log_softmax / logsumexp / logcumsumexp 全部采用"减最大值"策略；BCE logits 版本用 logsigmoid；SDPA 的 QK 点积强制提升到 float32 精度。
3. 与 PyTorch 对齐但保留差异：命名与默认参数大量参考 PyTorch（correction=1、count_include_pad、is_causal等），同时明确标注差异点（池化/卷积不限 2D、sparse_categorical_crossentropy类别轴必须为 -1、topk不支持sorted_=False、dropout 仅在 TRAINING 下生效）。
4. JIT 友好：masked_select/nonzero提供size定长参数以兼容 JIT 编译，sort等算子通过补齐到 2 的幂实现确定性的双调排序。

无论是搭建 CNN（conv2d+max_pool2d+avg_pool2d）、Transformer（scaled_dot_product_attention+layernorm+cross_entropy），还是做数值算法研究（qr/svd/einsum），本文覆盖的算子族都是最常用的工具箱；配合源码 docstring 中的可执行示例（文档站通过 markdown-exec 插件真实运行并输出结果），可以直接复制到项目中验证。

【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad