numpy中的张量收缩

问题描述

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

解决方法

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

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,466]])

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

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

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

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

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

添加

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

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

最后挤出1维：

In [186]: np.matmul( a[:,2))[:,:]
Out[186]: 
array([[ 40,466]])

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

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