numpy 中的张量收缩

Question

我有以下形状的 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) 似乎不是很有效。

Answer 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/@ 以将此 'outer' 行为更改为更常见的 'batch'。但是我们需要向 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)

numpy 中的张量收缩

Tensor contractions in numpy

python

numpy

numpy-ndarray