如何仅用Numpy操作无循环计算同数组嵌套函数求和的数学表达式

你当前的代码等价于计算所有i、j∈[0,N-1]的∏_{k=0}^{S-1} f(x_i[k], x_j[k])的总和，也就是双重求和嵌套沿S维度的乘积运算。

最优低内存实现方案 #

完全不需要生成(N,N,S)的超大数组，仅用(N,N)的临时数组+不超过20次的S维度循环即可实现，性能远高于原N维度循环方案：

import numpy as np

N, S = x.shape
# 初始化乘积矩阵为全1，仅分配(N,N)大小内存
prod_mat = np.ones((N, N), dtype=x.dtype)

# 仅循环S次（S<20，开销可忽略）
for k in range(S):
    col = x[:, k, None]  # 形状变为(N,1)，方便广播
    prod_mat *= f(col, col.T)

# 对整个矩阵求和得到最终结果
s = prod_mat.sum()

方案优势：

• 内存占用仅为(N,N)数组大小，按N=32768、float64类型计算仅占约8GB，符合你的内存要求
• 循环次数从3万次降到最多20次，性能提升至少三个数量级
• 兼容所有支持广播的对称f函数，通用性强

关于np.einsum的适配说明 #

np.einsum本身仅支持张量的线性乘积、缩并、转置操作，无法直接处理你用到的非线性运算（绝对值、最大值等），因此无法通过einsum一步直接替换所有循环实现需求。
如果你的f函数可以拆解为多项式形式，比如f(x,y) = (x-y)² = x² - 2xy + y²，可以进一步将乘积展开后用einsum计算，完全避免分配(N,N)数组，内存占用可以压到仅和N、S相关，适合内存更小的场景。

适配你提到的三类f函数的实现 #

仅需要替换上述循环中prod_mat *= f(col, col.T)一行即可：

• f为|x-y|：prod_mat *= np.abs(col - col.T)
• f为|x-y|²：prod_mat *= np.square(col - col.T)
• f为2 - max(x,y)：prod_mat *= 2 - np.maximum(col, col.T)

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