einsum 等价于 ndarray 乘法

Question

我有以下乘法例程：

import numpy as np
a = np.random.rand(3,3)
b = np.random.rand(3,50,50)
res = np.zeros((3, 50, 50))
for i in range(50):
    for j in range(50):
        res[:,i,j] = a @ b[:,i,j]

einsum 等价表达式是什么？

此致

Answer 1

可能想复习一下 Einstein summation notation:

res = np.einsum('ij, jkl -> ikl', a, b)

在这种情况下，np.tensordot也很有用：

np.tensordot(a ,b, 1).shape
Out[]: (3, 50, 50)