如何用单个np.einsum语句等价实现y=(x*A)@A.T计算？

要将分步的np.einsum调用合并为单个语句，完成y=(x*A)@A.T的计算（其中x为一维数组，A为二维矩阵，实际场景中A可能是1000x3或更大规模），可以按以下方式实现：

原始示例代码 #

import numpy as np
x = np.array([1.,2.,3.])
A = np.array([[3.,4.,5],
               [2.,2.,2.]])
y = (x*A)@A.T

# 需要替换占位符，写出单个einsum语句
# y2 = np.einsum('what to write here', x, A)

分步实现代码（参考） #

xA = np.einsum('i, ji->ji', x,A) # 对应逐元素相乘x*A
At = np.einsum('ij->ji', A) # 对应A的转置A.T
y2 = np.einsum('ij, jk', xA, At) # 对应矩阵乘法xA @ At
assert (y==y2).all() # 验证结果一致

合并后的单个np.einsum语句 #

直接使用以下语句即可完成计算，无需分步操作：

y2 = np.einsum('i, ji, ki->jk', x, A, A)

维度逻辑解释 #

• i：对应一维数组x的维度（示例中长度为3）
• ji：对应矩阵A的维度（j为行索引，i为列索引，示例中A是2行3列）
• ki：对应矩阵A的另一个引用，用来表示A.T的列行关系（因为A.T的[i,k]元素等于A的[k,i]元素）
• ->jk：指定输出结果的维度，对应最终m×m的矩阵（j和k均为原A的行索引，示例中结果为2行2列）

本质是对维度i求和，计算sum_i(x[i] * A[j,i] * A[k,i])，完全等价于(x*A)@A.T的计算逻辑。

验证代码 #

运行以下代码可确认结果一致：

assert (y == y2).all()
print("结果一致")

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