高维张量计算的几何本质：从shape、stride到广播与归约

ops-nn 这个名字，最早是我们内部给一个高维张量算子库起的代号。它要解决的事情，简单说就是：在神经网络训练和推理过程中，把各种形状乱七八糟的高维张量，高效地完成搬运、变形、合并和计算。叫它"几何引擎"，是因为我们后来发现，张量运算的本质，就是高维空间里的几何变换——形状是坐标系，步长是度量，逐元素算子是局部形变，矩阵乘法是线性映射，归约是沿着某个方向把空间压扁。把这些事想透了，写出来的算子不仅能跑得快，还更容易调试、更容易扩展到新的硬件上。

这套东西适合谁看？如果你写过几行神经网络代码，但是对"为什么 CPU 上转置一下就会慢一倍""为什么小 batch 也能把显存吃光""为什么同样的矩阵乘法，换一种循环顺序能快几十倍"这类问题感兴趣，这篇文章就是给你准备的。我会从最底层的设计思路开始，讲到 shape、stride、view、广播、归约这些基本操作到底在做什么，然后手写一个迷你张量计算核心，最后把我踩过的坑和排查方法整理成清单。全程不依赖任何框架内部的黑魔法，只谈原理和实操。

1. 从"循环"到"几何"：ops-nn 的底层设计哲学

1.1 朴素实现为什么慢？三层拆解

很多人第一次给张量库加算子，直觉就是写嵌套循环：shape 是几维就套几层 for，里层做加减乘除。这种实现能跑，但性能差得离谱。我实测过一次，在普通桌面级 CPU 上，用三层 Python 循环写一个 2D 卷积的前向计算，处理一张 224x224 的图、3x3 卷积核，跑完一次要几秒钟；而同样的逻辑用向量化方式实现，只要几十毫秒。差距不是两倍三倍，是几十倍上百倍。

差距来自三层：

第一层是循环本身的解释开销。Python 解释器逐条执行字节码，每一条算术都要经过对象系统，循环体里的每个元素访问都是函数调用级别的开销。这个层面解决靠的是批量操作，也就是把循环下沉到 C 层或者编译后的代码里。

第二层是内存访问模式。循环顺序不好，CPU 缓存命中率会非常低。比如按行存储的二维数组，如果循环内层先按列访问，每个元素都要重新加载缓存行，相当于人频繁在不同楼层之间跑上跑下。这个问题只靠向量化还不一定解决，还需要重新组织循环次序、分块访问。

第三层是计算资源没用满。现代 CPU 有 SIMD 指令，可以一次处理 4 个、8 个甚至 16 个浮点数；多数机器还有多个物理核心。朴素循环把这些资源全浪费了。所以 ops-nn 的第一个设计原则就是：能向量化就向量化，能并行就并行，绝不写裸循环。

1.2 张量的几何本质：一切运算都是空间变换

如果把一个张量看成高维空间里的一个"数量场"，很多计算机制就变得一目了然。

一个标量是一个点。一个向量是一条带方向的线，在几何上是一维空间里的位移。一个矩阵是一个二维数组，它描述的是从某个空间到另一个空间的线性映射，或者说是一组基到另一组基的变换关系。三维及以上的张量，可以理解为多个线性变换的"张量积"——把低维空间组合成高维空间。

这样一来，常见的张量操作就有了几何解释：

• 逐元素加法、乘法：对空间里的每一个点做局部变换，不改变空间结构。
• reshape：改变坐标系的划分方式，比如一个 2x3 的矩阵被拉成 6 个元素的向量，相当于把二维坐标重新编码成一维坐标。
• 转置/permute：交换坐标轴，空间结构没变，但轴的顺序变了，对应的"度量"（步长）要跟着调整。
• 广播：把一个低维空间的张量嵌入到高维空间里，每个位置放一份副本，本质是某种外积展开。
• 归约（sum、max、mean）：沿着某个轴把空间"压扁"，相当于高维空间投影到低维子空间。

"几何引擎"这个词就是这么来的。ops-nn 的算子，不管看起来多复杂，底层都是由这几类几何操作组合而成。理解了这一层，设计 API 的时候你会自然地倾向于"有视图就做视图，能惰性合并就合并"，而不是傻乎乎地在每个算子边界拷贝一次数据。

注意：几何视角不是文学比喻，它直接对应到数据结构上。张量的 shape 定义了空间有哪几个轴，stride 定义了每个轴走一步需要跳过多少内存单元。这两个字段一旦定了，所有运算的内存访问模式就定了，性能问题也由此而来。

2. 高维张量的核心操作机制：从形状推导到内存布局

2.1 shape、stride 与 view：视图机制到底在做什么

先看最底层的表示。一个密集张量在内存里通常是一段连续的一维数组，靠三元组来描述：

• shape：每一维的长度，比如 (2, 3, 224, 224)。
• stride：每一维相邻元素之间相差的元素个数。
• offset：起始位置的偏移。

对于一个连续存储的 NCHW 张量，shape 如果是 (N, C, H, W)，那么 stride 从外到内依次是 (CHW, H*W, W, 1)。也就是说，沿着 N 轴走一步，要跳过一整幅图的元素；沿着 W 轴走一步，只需要跳一个元素。

这个模型非常接近硬件实际。CPU 和 GPU 的内存都是线性的，多维索引全靠 stride 换算。ops-nn 里处理任何算子，第一步都是把高维索引映射到一维偏移，公式很简单：

offset = sum(idx[i] * stride[i]) + base_offset

有了这个公式，permute 和 transpose 就只是改 stride 和 shape，不动任何数据。我举个例子。一个 NCHW 连续张量经过 permute(0, 2, 3, 1) 变成 NHWC：

• 新 shape：(N, H, W, C)
• 新 stride：(CHW, W, 1, H*W)
  • N 轴还是跨一整幅图
  • H 轴只需要跨 W 个元素
  • W 轴跨 1 个元素
  • C 轴反而要跨 H*W 个元素

数据没动，但"每个轴在内存里跳多大步子"换了。这就是 view（视图）的实质：不复制，只改元数据。

在 Python 风格的代码里，这个逻辑写出来长这样：

class Tensor: def __init__(self, data, shape, stride=None, offset=0): self.data = data # 一维列表，模拟连续的底层内存 self.shape = tuple(shape) if stride is None: self.stride = self._default_stride(shape) else: self.stride = tuple(stride) self.offset = offset def _default_stride(self, shape): # 连续存储：C-contiguous 的默认步长 strides = [] acc = 1 for dim in reversed(shape): strides.append(acc) acc *= dim return tuple(reversed(strides)) def permute(self, *axes): new_shape = tuple(self.shape[ax] for ax in axes) new_stride = tuple(self.stride[ax] for ax in axes) return Tensor(self.data, new_shape, new_stride, self.offset)

这正是主流框架里"视图操作基本零开销"的秘密。如果你的自定义算子返回的是新拷贝，每次 permute 都要 O(N) 地复制数据；如果返回视图，耗时几乎可以忽略。我在 ops-nn 里定了一条规则：

凡是只调整形状、轴顺序、扩展维度的操作，默认返回视图。只有当你显式调用 copy/contiguous 时才真正复制数据。

2.2 广播（broadcast）背后的空间嵌入逻辑

广播是让新手最头疼的机制，但几何视角下它非常自然。想象一个 shape 为 (3,) 的向量和一个 shape 为 (4,) 的向量相加，直接相加是不合法的。可是如果第一个向量先"嵌入"到一个 (1, 3) 的行向量，第二个向量变成 (4, 1) 的列向量，再按外积的方式扩展成 (4, 3) 的公共空间，两者就可以相加了。

这就是广播的本质：把一个低维张量，沿着缺失的轴和长度为 1 的轴，逻辑上复制多次，嵌入到高维空间里。

标准广播规则只有三条：

• 从尾部对齐 shape。
• dim 长度为 1 的轴可以扩展到目标长度。
• 缺失的维度默认长度为 1，然后按前一条扩展。

例如：

shape A: (3, 1) shape B: (4) 对齐后 B: (1, 4) 扩展公共形状: (3, 4)

从几何上，A 沿第二维"压扁"成 1 之后被拉伸到 4，B 沿第一维被拉伸到 3。最终每个元素都是"两个输入对应空间坐标的某种组合"。

实现广播时，ops-nn 会算出一个"广播后的共同 shape"，但不实际生成副本。内存访问还是通过 stride 实现：对长度为 1 的轴，stride 设成 0 即可。这样在循环里访问 A[i][j] 时，如果 A 的某个轴 stride 是 0，它永远指向同一个元素，等效于复制但零内存开销。

这带来一个非常实用的技巧：判断两个张量能不能相加，不用去记广播规则的表格，用手里的 shape 从右往左逐位比较就行，那位是 1 就能扩张，不是 1 就得完全一样。

但我必须提醒一个反直觉的坑。广播对"加法、逐元素乘法"这种局部运算是直观的，但对"归约+广播"组合就容易出错。比如你想用(N, 1)的均值去归一化(N, D)的特征，x / mean是合理的；但如果你顺手把 mean 的 keepdims 改成 False，均值就变成(N,)，结果会变成(N, D)除以(N,)的对齐广播——大概率错了，还不会报错。这个问题我在 4.1 节会展开说。

2.3 归约与规约的"压低一维"效应

归约操作（sum、mean、max、min）本质上是沿着某一个轴做聚合，把那个方向上的信息压成一条线。几何上这很像投影：三维物体投影到二维平面，丢掉一个维度。

实现上的关键在于：归约轴在内存里访问的连续性。以 (N, C, H, W) 为例：

• 对 C 轴做 sum，内存访问模式是每隔 H*W 个元素取一个贡献，步长很大。
• 对 W 轴做 sum，内存访问是连续往后加，全部命中缓存。

同样的计算量，沿 W 归约通常比沿 C 归约快一个量级。这是我在优化算子时碰到的最常见的性能差异来源。

归约的数值问题也比想象中严重。几个大数相加，小数会悄悄被吞掉，这在高维张量归约里特别常见。比如在 logits 上做 softmax，如果直接对 exp 后的结果求和，大 logit 会导致数值溢出；所以稳定做法是先减掉最大值再算 exp。对于一般的 sum，也可以用分段求和来降低误差。

ops-nn 里的分组归约实现可以这样理解。假设要把沿最后一维归约，最直接的写法是：

# 朴素版本：嵌套循环，逐元素累积 for n in range(N): for c in range(C): acc = 0.0 for i in range(H): for j in range(W): acc += x[n, c, i, j] out[n, c] = acc

向量化版本会把它拆成"扁平化循环 + 两两相加"：

# 向量化思路：先把输入理解为 (N*C, H*W)，再对 H*W 这层做树状求和 flat = x.reshape(N * C, H * W) out = flat.sum_axis_vectorized(axis=1) # 内部使用 SIMD

这既提升了吞吐，也在一定程度上有益于数值稳定性，因为树状求和的误差累积比顺序求和小。

3. 实操剖析：亲手写一个迷你张量计算核心

3.1 数据结构与接口设计

这一节我会带你写一个能跑、能测、能扩展的迷你张量库，名字沿用 ops-nn 的精神。先定义最底层的 Tensor 类，我刻意让它足够简单，方便看清每一步。

class Tensor: """迷你张量核心：连续内存 + shape + stride + offset。""" def __init__(self, data, shape, stride=None, offset=0): self.data = data self.shape = tuple(shape) if stride is None: stride = [] acc = 1 for d in reversed(shape): stride.append(acc) acc *= d stride = list(reversed(stride)) self.stride = tuple(stride) self.offset = offset @property def ndim(self): return len(self.shape) @property def size(self): acc = 1 for d in self.shape: acc *= d return acc def _idx_to_offset(self, idx): off = self.offset for i, s in zip(idx, self.stride): off += i * s return off def __getitem__(self, idx): if not isinstance(idx, tuple): idx = (idx,) return self.data[self._idx_to_offset(idx)] def reshape(self, *new_shape): # 只在元素总数相同时允许，且返回视图 assert self.size == self._prod(new_shape), "shape incompatible" return Tensor(self.data, new_shape, None, self.offset) def permute(self, *axes): assert sorted(axes) == list(range(self.ndim)) new_shape = tuple(self.shape[ax] for ax in axes) new_stride = tuple(self.stride[ax] for ax in axes) return Tensor(self.data, new_shape, new_stride, self.offset) def broadcast_to(self, target_shape): # 从尾部对齐广播，stride 中长度 1 的轴设成 0 pad_ndim = len(target_shape) - self.ndim new_stride = [0] * pad_ndim + list(self.stride) new_shape = [1] * pad_ndim + list(self.shape) final_stride = [] for i in range(len(target_shape)): if new_shape[i] == 1: final_stride.append(0) else: assert new_shape[i] == target_shape[i] final_stride.append(new_stride[i]) return Tensor(self.data, target_shape, final_stride, self.offset) @staticmethod def _prod(shape): acc = 1 for d in shape: acc *= d return acc

接口设计的几个细节，我想单独说。

第一，reshape默认返回视图，这意味着它不检查数据是否真的连续。如果对一个转置过的张量做 reshape，新的 shape 和 stride 对不上，后续运算可能产生完全错误的结果。为了安全，实际实现的 reshape 应该先判断is_contiguous，非连续时悄悄拷贝，或者直接抛错提示用户调用contiguous()。实操中我建议在 debug 模式下做一次校验。

第二，broadcast_to返回视图而不是真实展开。stride 为 0 的轴意味着"每次访问同一块地址"，这保证了广播不出内存开销。整套设计的目标就是：能共享底层数据就绝不多搬一次。

3.2 核心算子的向量化与内存访问优化

有了 Tensor，下一步是写算子。我拿三个最有代表性的场景示范：逐元素运算、矩阵乘法、卷积的矩阵化视角。

逐元素运算（比如y = alpha * x + beta）是非常典型的 memory-bound 任务。它的瓶颈不在计算，而在把数据从内存搬到寄存器。优化手段是：先扁平化，再按连续内存块遍历，同时用 SIMD 批量处理：

def add_scalar(x: Tensor, alpha: float) -> Tensor: out = Tensor([0.0] * x.size, x.shape, None, 0) # 实际内核里这里会走向量化：一次处理多个元素 for i in range(x.size): out.data[i] = x.data[x.offset + i * x.stride[0]] + alpha return out

但要注意，如果 x 是非连续的（比如转置过的视图），上面这个扁平化循环就错了，因为最后一个轴的步长不一定是 1。处理非连续张量的通用办法是写一个"通用遍历"版本，它通过递归或者多维索引来访问：

def apply_unary(x: Tensor, fn): out = [0.0] * x.size def rec(idx, dim): if dim == x.ndim: out[x._idx_to_offset(idx)] = fn(x[idx]) return for i in range(x.shape[dim]): rec(idx + (i,), dim + 1) rec((), 0) return Tensor(out, x.shape, None, 0)

通用版本覆盖所有情况，但性能不如连续版本。所以真正的算子里通常写两条路径：is_contiguous走快速路径，否则走通用路径。这也是主流框架的常见做法。

矩阵乘法是 compute-bound 的典型，优化它我只会强调分块。朴素三重循环是 i、j、k 顺序，内层 k 循环每次移动两个矩阵的行/列，缓存命中率很差。分块的做法是把 A 和 B 都切成小块，保证一个块能完整放入 L2 缓存，然后块内做密集乘加：

def matmul(A: Tensor, B: Tensor, block_size=32): M, K = A.shape K2, N = B.shape assert K == K2 C = [0.0] * (M * N) # 外层分块：i 和 j 两个块 for i0 in range(0, M, block_size): for j0 in range(0, N, block_size): for k0 in range(0, K, block_size): for i in range(i0, min(i0 + block_size, M)): for j in range(j0, min(j0 + block_size, N)): acc = 0.0 a_row_base = i * K b_col_base = j for k in range(k0, min(k0 + block_size, K)): acc += A.data[a_row_base + k] * B.data[b_col_base + k * N] C[i * N + j] += acc return Tensor(C, (M, N), None, 0)

这段代码的关键是循环顺序。把 k 的循环放在最内层，并且访问 B 时按列跳跃——对于按行存储的 B，这其实不是最优的。更高效的做法是先把 B 转置成按行访问友好的布局，也就是"打包"。我实操时的经验是：

• 对按行优先存储的矩阵，A 按行访问没问题，B 需要转置或重排，让内层乘积累加到连续内存。
• 块大小不要贪大。我试过 32x32 在多数台式 CPU 上效果最好，64x64 在某些机器上反而下降，因为缓存放不下。

卷积的矩阵化视角也值得一提。把输入图像通过 im2col 提取成一个大矩阵，每个卷积窗口拉成一行/列，卷积就变成了矩阵乘法。代价是内存占用变为原来的 kernel_size 平方倍。如果你的环境内存紧张，用这个方案要谨慎。而 Winograd 或 FFT 则是在数学上减少乘法次数，但会增加变换计算的复杂度。真正的工程选择要看实际部署硬件，我在 5.3 节会再说。

3.3 算子融合的实战收益：一个例子

算子融合是个经常被提起的概念，但在小型自研库里同样能落地。举个最常见的例子：计算y = a * x + b。如果拆成乘法和加法两步，每个算子都会完整遍历一次数据、写一次中间结果、再读一次。内存带宽被吃掉两次。融合后，单次遍历同时完成乘加：

def fused_axpy(x: Tensor, a: float, b: float) -> Tensor: out = [0.0] * x.size # 假设 x 连续 base = x.offset stride0 = x.stride[0] if x.ndim == 1 else 1 for i in range(x.size): val = x.data[base + i * stride0] out[i] = a * val + b return Tensor(out, x.shape, None, 0)

如果 x 是多维的、非连续的，可以先用通用遍历的框架做融合，逻辑一样：在同一个循环体内完成乘加。这个思想在优化大批量归一化、激活层、dropout 组合时会带来肉眼可见的加速。我在 benchmark 里测过，融合后局部运算的耗时通常能降 30%~50%，因为省掉了中间张量的读写。

4. 常见问题与排查技巧实录

4.1 广播语义引发的隐性错误

我在 2.2 节提过这个坑，这里展开讲。典型的错误场景是：你有一批预测值 pred，shape 是 (N, C)，对应的 target shape 是 (N,)，想算每个样本的 MSE。于是你写：

loss = ((pred - target) ** 2).mean()

如果 target 忘了 reshape 成 (N, 1)，广播规则会把 (N,) 对齐到 (N, C)：target 相当于沿 C 维复制 C 份。如果 C 恰好等于 N，这个错误不会报错，但 loss 的含义已经完全错了，而且模型还能正常训练下去，只是效果变差。这属于最难排查的静默错误。

我的排查习惯有三条：

• 所有关键运算前，用断言检查 shape 是否完全符合预期，或者打印交接处的 shape。debug 模式下的 shape 断言成本很低，但能救回大量调参时间。
• 写一层极薄的"形状契约"装饰器，在每个算子入口自动核对输入 shape 和输出 shape 的推导关系。这在小库维护里非常实用。
• 测试用例必须覆盖"维度不匹配但能广播"的场景，明确告诉你这个错误会被静默吞掉，然后设计成直接抛异常。

4.2 视图与拷贝的脏数据问题

视图机制给你性能，也给你一堆坑。最典型的是就地操作和视图共享内存。假设你写：

a = tensor_from_list([[1, 2], [3, 4]]) b = a.permute(1, 0) # 转置视图 b[0, 0] = 99 # 修改 b 的一个元素 print(a[0, 0]) # a 也变了

从设计上，视图就是共享数据，这个行为完全正确。但初用者很容易被"看起来是不同变量"骗了。我处理这个问题的方式是：在 Tensor 上给每个视图增加一个base字段，指向最原始的底层缓冲区；任何会触发就地写的操作，先检查是否与其他张量共享数据，必要时自动copy()。这会有性能损耗，所以只对写入操作做检查，读操作完全放开。

另一个更隐蔽的坑是"非连续视图上做隐形拷贝"。有些算子（比如某些 norm 或 einsum 实现）强行要求连续内存，遇到非连续输入会默默拷贝。你以为是 O(1) 的操作，实际变成 O(N) 拷贝，性能瞬间崩掉。排查这类问题没有捷径，只能靠 profiler。我建议在自定义算子的关键路径上，显式调用类似make_contiguous的逻辑，并注释掉"这里的拷贝是故意的还是被迫的"。如果被迫，就回头研究能不能改成支持 stride 访问的通用实现。

4.3 浮点累积与性能陷阱

浮点误差在高维张量里会被放大，尤其是在大批量归约和混合精度场景。

一个印象很深的例子是：我在做某个跨平台系统的 mean 算子时，发现同一批输入在不同 batch size 下输出的均值不一样，差异最大到小数点后第三位。原因很简单：当 batch 很大时，直接累加 10 万个浮点数，最后加进去的小数字往往会被前面的大数"吃掉"。解决方案在标准库设计里已经存在：用两趟法或者分组求和。两趟法先算大致的均值，再算每个元素与均值的差，最后修正。这样额外增加一次遍历，但稳定性提升明显。

性能陷阱方面，最常见的两个：

• 对非连续张量做逐元素运算，没有走快速路径，导致每个元素都要通过步长计算偏移。我见过一个算子因此比连续版本慢了 6 倍。
• 频繁创建形状相同的小张量，造成碎片化内存申请。解决方式是在库里加一个简单的内存池，或者复用临时缓冲区池。框架级的"缓存分配器"就是这么来的。

4.4 数值梯度检查与调试技巧

要验证一个自定义算子对不对，最可靠的是梯度检查。朴素原理是：对输入加一个很小的扰动，看输出的变化是否等于解析梯度乘以扰动。

def check_gradient(fn, x, analytic_grad, eps=1e-6): # 逐元素扰动 numerical_grad = [] for i in range(x.size): x_plus = x.data[:] x_minus = x.data[:] x_plus[i] += eps x_minus[i] -= eps fx_p = fn(Tensor(x_plus, x.shape)).data[0] fx_m = fn(Tensor(x_minus, x.shape)).data[0] numerical_grad.append((fx_p - fx_m) / (2 * eps)) return numerical_grad, analytic_grad

实操中我还有两个技巧：

• 用小 shape 测试，比如 (2, 3)，而不是 (256, 256, 3)。形状大了跑得慢，且手算参照麻烦。
• 用相对误差而不是绝对误差。浮点截断会让绝对误差在数值很小的地方看起来很大。相对误差小于 1e-5 基本就放心了。

调试高维张量的另一个痛点是打印不方便。我给 Tensor 写了一个摘要打印函数：超过某个阈值就显示 shape、stride、前几个元素和内存地址片段。肉眼盯着 224x224 的矩阵看，什么都看不出来；看 shape、stride、首地址，很多问题立即暴露。

5. 同类型方案对比与落地选型建议

5.1 动态图与静态图的张量内核差异

在很多自研项目和论文复现场景里，要不要引入某种计算引擎，往往要在动态图风格和静态图风格之间做选择。

动态图风格的优点是调试直观：算子立即执行，shape 信息随手可见。缺点是每次迭代都有解析和分发的开销，遇到大量小算子时性能折损明显。静态图风格正好反过来：先构造计算图，再统一编译执行，能做算子融合和内存复用，但调试体验相对差，中途想看某个中间结果要额外标记。

ops-nn 在早期选择了类似动态图逐算子执行的路线，但我发现大量小算子叠加后，整体运行时间被"内层循环+函数调用"拖累了。后来我在设计中加了一个"脚本化"入口：记录一串算子调用，然后交给一个迷你 planner 按顺序做融合打包。这不算完整的静态图编译，但已经能解决 80% 的性能问题。具体融合策略就是 3.3 节说的：相邻的逐元素运算合并成一个 pass，归约和矩阵乘尽量保持独立以利用专用内核。

5.2 衡量一个算子库的关键维度

挑库或者自己写库的时候，我一般用几个维度横向评估：

我见过不少人只盯着性能选型，结果遇到一个奇怪的广播 bug，光排查就花了两天。所以我现在选型时，"可调试性"的权重会拉到和性能一样高。

5.3 我们最终采用的实现风格和取舍

综合前面的经验，我目前在做一个自研推理框架时，张量内核部分的实现风格可以总结为几条：

• 核心数据结构保留shape + stride + offset + base四件套，所有算子优先走视图路线。
• 每个算子实现两个入口：连续快速路径和通用路径。快速路径要求输入连续且 dtype 一致；通用路径通过步长映射支持任意视图。快速路径如果没实现，宁可抛异常提示，也不要静默走慢路径。
• 逐元素小算子默认融合。相邻的 add、mul、激活、scale 会在调度层打包成一个 kernel，减少中间读写。
• 矩阵乘法等 compute-bound 算子用分块加打包，块大小和目标缓存参数化，方便针对不同机器调优。
• 所有归约实现都加 keepdims 后的通用化处理，并且用分组求和缓解累积误差。

这套方案在 IPC 和内存带宽的实测上，能赶上主流框架的 70%~90% 水平（当然我只测了 CPU 端，GPU 端集群版是另一个团队在维护）。作为自研项目或教学剖析的骨架，这个性价比已经非常高。

最后分享一个小经验

如果你看完这篇想自己动手写一个张量核心，我最大的建议是：先不要急着仿照某个成熟框架抄 API，而是从 shape、stride、offset 这三个字段开始，亲手实现 reshape、permute、broadcast、sum 这四个算子。你能走通这四个操作，基本上就理解了高维张量计算的全部关键机制：空间划分、视图共享、内存访问模式和归约的数学生态。之后再往上加 matmul、conv、softmax，都只是"分块优化"和"数值稳定"的重复应用。

我在捣鼓 ops-nn 时最有成就感的一件事，不是某个算子跑得比主流框架还快，而是终于明白了一个道理：张量库的代码量其实不大，难的是把每一步内存访问都放在你脑子里的那张"几何地图"上。把高维形状看成空间、把步长看成度量，所有算子的设计逻辑都会变得透明起来。这套视角，希望你也能在实操里感受到。