np.tensordot 的核心是 axes 参数指定的轴对收缩操作,必须为等长元组,输出形状为未收缩轴按 a 前 b 后拼接;不支持轴重排或批处理,性能优于 einsum 但灵活性低,需注意 dtype、内存连续性及多轴收缩效率。
np.tensordot 不是通用“张量乘法”,而是按指定轴对齐后求和的收缩操作。关键在 axes 参数:它必须是长度相等的两个元组,分别对应第一个和第二个数组要收缩的轴索引。比如 axes= 表示把 
a 的第 1、2 轴与 b 的第 0、1 轴逐对配对并求和。
常见错误是传入单个整数(如 axes=1)——这仅适用于两数组最后/最前各 1 维自动匹配的简化场景,一旦维度不规则或需跨位置收缩,必须显式写成元组形式。
a.shape == (2, 3, 4),b.shape == (3, 4, 5),想收缩中间两维,用 np.tensordot(a, b, axes=([1, 2], [0, 1])),结果 shape 是 (2, 5)
axes=2,numpy 会尝试收缩最后 2 维,但 a 最后两维是 (3, 4),b 最后两维是 (4, 5),尺寸不匹配直接报 ValueError: shape mismatch for sum
axes=([3], [0]) 在 3 维数组上会触发 AxisError
很多人想用 np.tensordot 替代 np.einsum,但二者能力不对等。np.tensordot 只能做「收缩 + 拼接」,不能重排输出轴顺序,也不能保留未收缩轴的任意排列。例如:np.einsum('ijk,kl->ijl', a, b) 中输出轴顺序是 i-j-l,而 np.tensordot 的输出固定为「未收缩的 a 轴 + 未收缩的 b 轴」顺序。
性能上,tensordot 底层调用高度优化的 BLAS,当收缩涉及大量连续内存访问(如矩阵乘、批量矩阵乘)时,比 einsum 快;但一旦需要复杂索引逻辑,强行用 tensordot 可能得先 transpose,反而更慢且易错。
a.shape == (B, I, J),b.shape == (B, J, K),想做每组 B 内的矩阵乘,应改用 np.einsum('bij,bjk->bik', a, b) 或 np.matmul;tensordot 无法天然表达这种“按批独立收缩”tensordot 对 float32 / float64 友好,但对 int64 输入可能因累加溢出产生静默错误,建议显式转 float32 再算假设你有两个 4D 张量 a 和 b,shape 都是 (N, C, H, W),想对 C 和 H 维度做内积(即通道与空间联合压缩),保留 N 和 W,可这样写:
result = np.tensordot(a, b, axes=([1, 2], [1, 2]))
注意:结果 shape 是 (N, W, N, W),因为未收缩轴是 a 的 [0, 3] 和 b 的 [0, 3],拼起来就是 (N, W, N, W)。如果你只想要 (N, N)(比如 batch-wise 相似度),就得额外再收缩 W 维:
result = np.tensordot(a, b, axes=([1, 2], [1, 2])) → (N, W, N, W)
final = np.tensordot(result, np.ones(W), axes=([1, 3], [0, 0])),或更直白地 final = result.sum(axis=(1, 3))
np.einsum('nchw,nchw->nn', a, b),语义清晰且不易错np.tensordot 默认输出 dtype 由输入决定,但内部计算可能升格(如 int32 × int32 → int64),尤其在累加维度很大时。如果后续要 GPU 加速或内存敏感,务必检查 result.dtype 是否符合预期。
另一个隐形坑是内存连续性:tensordot 对非 C-contiguous 数组(比如由 transpose 或切片生成)会先复制为连续内存,大幅拖慢速度。遇到性能差时,先用 a = np.ascontiguousarray(a) 显式转换。
a.flags.c_contiguous
np.tensordot(a.transpose(0, 2, 1), b, ...) → 先 a_t = np.ascontiguousarray(a.transpose(0, 2, 1)) 再用