numpy中的张量收缩答案

【问题标题】：Tensor contractions in numpynumpy中的张量收缩
【发布时间】：2021-07-11 10:26:59
【问题描述】：

我有 numpy 形状的数组：(2000, 3) (3, 2000, 2)。我想知道以下张量收缩是否：

t1 = np.einsum("ij,jik->ik", a, b)

可以使用np.dot 完成。我已经尝试过：

  out = np.zeros((2000, 2))
  for k in range(2):
      out[:, k] = np.dot(a, b[:, :, k])

但它会引发错误： ValueError: 无法将输入数组从形状 (1000,1000) 广播到形状 (1000)

迭代第一个维度 (i) 似乎效率不高。

【问题讨论】：

在i, 2000, since the dot 上迭代` 产品总和在j 上。 k 未共享。 matmul 可以在没有循环的情况下完成，但需要转置和维度添加。

标签： python numpy numpy-ndarray

【解决方案1】：

有一个最小的例子（我们喜欢在问题中看到）：

In [173]: a = np.arange(12).reshape(4,3); b=np.arange(24).reshape(3,4,2)
In [174]: t1 = np.einsum('ij,jik->ik', a, b)
In [175]: t1.shape
Out[175]: (4, 2)
In [176]: t1
Out[176]: 
array([[ 40,  43],
       [136, 148],
       [268, 289],
       [436, 466]])

工作循环，i 上的迭代（是的，即使它可能很大）

In [177]: out = np.zeros((4,2),int)
In [178]: for i in range(4):
     ...:     out[i,:] = np.dot(a[i], b[:,i,:])
     ...: 
In [179]: out
Out[179]: 
array([[ 40,  43],
       [136, 148],
       [268, 289],
       [436, 466]])

在我们担心效率之前，它必须工作！

现在，如果我们将b 转置，我们可以将einsum 更改为：

In [181]: np.einsum('ij,ijk->ik', a, b.transpose(1,0,2)).shape
Out[181]: (4, 2)

j 在正确的位置，dot 可以工作 - 但它在共享的 i 维度上执行 outer 产品。我们必须取它的对角线：

In [182]: np.dot( a, b.transpose(1,0,2)).shape
Out[182]: (4, 4, 2)

添加了matmul/@ 以将此“外部”行为更改为更常见的“批处理”。但我们需要为a 添加一个维度，因此i 是所有术语的3 个中的第一个：

In [184]: np.matmul( a[:,None,:], b.transpose(1,0,2)).shape
Out[184]: (4, 1, 2)

最后挤出1维：

In [186]: np.matmul( a[:,None,:], b.transpose(1,0,2))[:,0,:]
Out[186]: 
array([[ 40,  43],
       [136, 148],
       [268, 289],
       [436, 466]])

最后一个表达式应该和einsum 一样快，如果不是更快的话。

In [192]: timeit t1 = np.einsum('ij,jik->ik', a, b)
8.05 µs ± 63.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [193]: timeit np.matmul( a[:,None,:], b.transpose(1,0,2))[:,0,:]
5.19 µs ± 183 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

【讨论】：