如何将Numpy中的einsum操作替换为标准数组运算？

针对你给出的np.einsum('aijb,rjb->ira', a, b)，可以通过广播+求和或者reshape+矩阵乘法两种方式实现，以下是具体代码和原理说明：

维度对应关系梳理 #

• 输入a的维度：(a=128, i=16, j=8, b=32)
• 输入b的维度：(r=256, j=8, b=32)
• 输出output的维度：(r=256, i=16, a=128)
• 核心逻辑：对j和b两个维度做元素相乘后求和，最终调整维度顺序为(r,i,a)

方法1：广播+求和+转置 #

通过扩展维度实现广播对齐，再对目标维度求和，最后调整维度顺序：

import numpy as np
a = np.random.rand(128, 16, 8, 32)
b = np.random.rand(256, 8, 32)

# 扩展维度实现广播对齐
a_expanded = a[np.newaxis, :, :, :, :]  # shape: (1, 128, 16, 8, 32)
b_expanded = b[:, np.newaxis, np.newaxis, :, :]  # shape: (256, 1, 1, 8, 32)

# 元素相乘后对j、b维度求和
product = a_expanded * b_expanded
summed = product.sum(axis=(-1, -2))  # 对最后两个维度求和，shape: (256, 128, 16)

# 转置得到目标维度顺序
output1 = summed.transpose(0, 2, 1)  # shape: (256, 16, 128)

方法2：reshape+矩阵乘法+转置 #

将需要求和的j、b维度合并为一维，利用矩阵乘法的BLAS优化加速运算（适合大数组场景）：

# 合并j、b维度为一维
a_flat = a.reshape(128 * 16, 8 * 32)  # shape: (2048, 256)
b_flat = b.reshape(256, 8 * 32)  # shape: (256, 256)

# 矩阵乘法实现批量内积
matmul_result = b_flat @ a_flat.T  # shape: (256, 2048)

# 恢复维度并转置
output2 = matmul_result.reshape(256, 128, 16).transpose(0, 2, 1)  # shape: (256, 16, 128)

验证结果一致性 #

可以通过np.allclose验证两种方法的输出和原einsum结果是否一致：

output_einsum = np.einsum('aijb,rjb->ira', a, b)
print(np.allclose(output1, output_einsum))  # 输出True
print(np.allclose(output2, output_einsum))  # 输出True

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