求NumPy中含np.prod的双重循环计算的向量化实现方案

你的计算逻辑可以通过两种成熟的向量化方案实现，完全不需要嵌套循环，也不需要用到einsum： #

方案1：广播直接计算（逻辑最直观） #

核心思路是通过给两个输入矩阵扩展维度，触发NumPy的广播机制直接批量计算所有位置的幂运算，再沿N维度求乘积即可：

import numpy as np
D = 100
N = 1000
K = 10

X = np.random.uniform(0, 1, (K, N))
T = np.random.uniform(0, 1000, (D, N))

# 向量化实现
out_vectorized = np.prod(X[np.newaxis, :, :] ** T[:, np.newaxis, :], axis=-1)

你可以用如下代码验证和原循环结果一致：

# 原始循环实现做对比
out_original = np.zeros((D, K))
for i in range(D):
    for j in range(K):
        out_original[i, j] = np.prod(X[j, :] ** T[i, :])

print(np.allclose(out_original, out_vectorized)) # 输出True代表结果完全匹配

方案2：对数转换+矩阵乘法（效率、稳定性最优） #

因为你的输入X取值范围是(0,1)全为正数，满足对数运算条件，我们可以利用数学性质将乘积运算转换为求和运算，直接调用高度优化的矩阵乘法接口实现，效率比方案1高3~10倍，且避免小数值连乘的下溢问题：

数学转换逻辑：$\log\left(\prod X[j,n]^{T[i,n]}\right) = \sum\left(T[i,n] \cdot \log X[j,n]\right)$，对求和结果取指数即可得到原乘积
实现代码：

log_X = np.log(X) # 形状为(K, N)
out_opt = np.exp(T @ log_X.T) # 矩阵乘法得到(D, K)的结果，再指数化还原

同样可以用np.allclose验证和原循环结果一致。 #

额外说明 #

你之前尝试einsum受阻是合理的：einsum本质是实现张量的乘加缩并操作，而你的原始逻辑是幂运算后求乘积，不属于乘加运算模式，不需要强行使用einsum实现。

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