如何在NumPy中无额外内存分配实现数组相乘并沿轴求和？

架构老林

2026-7-25

问题

是否存在无需分配额外内存即可对两个数组执行乘法运算并沿单个或多个轴求和的方法？示例代码如下：

import numpy as np

A = np.random.random((10, 10, 10))
B = np.random.random((10, 10, 10))

C = np.sum(A[:, None, :, :, None] * B[None, :, None, :, :], axis=(-1,-2))

在计算C时，会先创建一个尺寸为10×10×10×10×10的中间矩阵，随后立即被压缩。请问在NumPy中能否避免这种情况？

解决方案

完全可以避免生成这类大内存开销的中间数组，NumPy提供了两种高效的原生方法来直接完成这种乘法+多轴求和的操作（本质是张量收缩）：

1. 使用爱因斯坦求和约定（`np.einsum`）

np.einsum通过索引符号直接描述张量间的运算关系，无需显式广播数组，内部会优化计算路径，避免中间大数组的生成。

针对你的示例，对应的einsum写法如下：

import numpy as np

A = np.random.random((10, 10, 10))
B = np.random.random((10, 10, 10))

# 按索引对应关系计算：对A的第3轴（l）和B的第2、3轴（l,m）求和
C_einsum = np.einsum('ikl,jlm->ijk', A, B)

2. 使用张量点积（`np.tensordot`）

np.tensordot专门用于张量间的收缩运算，通过指定待求和的轴对来完成计算，同样不会生成中间广播数组。

对应示例的写法如下：

# 指定A的第2轴与B的第1轴配对，同时对B的第2轴求和（等价于原操作的双轴求和）
C_tensordot = np.tensordot(A, B, axes=([2], [1])).sum(axis=-1)

验证正确性

可以通过以下代码验证两种方法的结果与原代码完全一致：

C_original = np.sum(A[:, None, :, :, None] * B[None, :, None, :, :], axis=(-1,-2))

print(np.allclose(C_einsum, C_original))  # 输出True
print(np.allclose(C_tensordot, C_original))  # 输出True