如何无需循环对(m,n,n)型NumPy数组的子数组应用函数？

假设你有两个NumPy数组：

• arr：形状为(m, n, n)的数组，包含m个n×n的子数组
• weights：形状为(n, n)的权重数组

要完成“对每个子数组逐元素乘权重，再求和得到长度m的数组”的需求，完全不需要用for循环，利用NumPy的广播机制或张量点积就能高效实现：

方法一：广播 + 指定轴求和 #

利用NumPy的广播规则，weights会自动扩展为(m, n, n)的形状，与arr逐元素相乘；之后对相乘结果的最后两个轴求和，直接得到每个子数组的运算结果：

import numpy as np

# 示例数据
m, n = 5, 3
arr = np.random.rand(m, n, n)
weights = np.random.rand(n, n)

# 核心运算
result = np.sum(arr * weights, axis=(1, 2))

方法二：张量点积（tensordot） #

逐元素相乘再求和本质是两个n×n数组的元素级点积，用np.tensordot可以直接指定维度进行运算：

result = np.tensordot(arr, weights, axes=((1, 2), (0, 1)))

验证结果 #

两种方法得到的结果完全一致，也和循环实现的结果等价：

# 循环实现（用于对比）
loop_result = np.array([np.sum(arr[i] * weights) for i in range(m)])
print(np.allclose(result, loop_result))  # 输出True，验证正确性

内容的提问来源于stack exchange，提问作者roshoka