numpys张量点的向量化计算答案

【问题标题】：Vectorized computation of numpys tensor dotnumpys张量点的向量化计算
【发布时间】：2020-12-17 19:24:24
【问题描述】：

我有两个向量，分别包含形状为 (3,3) 和形状为 (3,3,3,3) 的张量。向量具有相同的长度，我正在计算这两个向量的逐元素张量点。例如，想要对以下计算进行矢量化以提高性能：

a = np.arange(9.).reshape(3,3)
b = np.arange(81.).reshape(3,3,3,3)
c = np.tensordot(a,b)

a_vec = np.asanyarray([a,a])
b_vec = np.asanyarray([b,b])    
c_vec = np.empty(a_vec.shape)

for i in range(c_vec.shape[0]):
    c_vec[i, :, :] = np.tensordot(a_vec[i,:,:], b_vec[i,:,:,:,:])

print(np.allclose(c_vec[0], c))
# True

我考虑过使用numpy.einsum，但无法找出正确的下标。我尝试了很多不同的方法，但到目前为止都失败了：

# I am trying something like this
c_vec = np.einsum("ijk, ilmno -> ijo", a_vec, b_vec)

print(np.allclose(c_vec[0], c))
# False

但这并不能重现我上面想要的迭代计算。如果使用 einsum 无法做到这一点，或者有更高效的方法可以做到这一点，我愿意接受任何类型的解决方案。

【问题讨论】：

将输入数组加倍有什么意义？本部分：a_vec = np.asanyarray([a,a])
@Divakar 是的，我正在尝试对循环进行矢量化。我有很多张量，希望以这种方式加快速度。
@Marat 真的没有意义，我只是将其用作健全性检查。

标签： python arrays numpy vectorization tensordot

【解决方案1】：

np.einsum 的矢量化方式将是 -

c_vec = np.einsum('ijk,ijklm->ilm',a_vec,b_vec)

【讨论】：

【解决方案2】：

tensor_dot 有一个 axes 参数，您也可以使用：

c_vec = np.tensordot(a_vec, b_vec, axes=([1, 2], [1, 2]))

【讨论】：