Python 结合sparse和einsum执行大型稀疏和

我有一个形状=（N，N）的矩阵a和一个形状=（N，N）相同的矩阵B。 我正在使用以下einsum（使用库）构造矩阵M：

这是对以下总和的评估：

这使得矩阵M具有形状（N，N，N，N）。然后，我将其重塑为二维形状数组（N**2，N**2）

这必须是二维的，因为它被视为线性操作符

我想使用这个库，因为M是稀疏的，并且变得太大，无法存储大的N。我可以使A和B稀疏，并将它们插入到

oe.contract

中

---

问题是，稀疏仅支持2D数组，因此无法生成形状（N，N，N.N）的4D输出。有没有一种方法可以将einsum和重塑步骤结合起来，以允许以这种方式使用稀疏，因为M的最终形状是二维的？

这可能对您使用

opt_einsum

没有帮助，但通过一点重新组织，我可以大大加快

np.einsum

的速度，至少对于小型阵列是这样

做两个

B

的部分乘积：

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N,N)

pq

对是相同的，因此我们不需要重新计算：

c2 = np.einsum('nm,onm,pnm->op',A,c1,c1)

我验证了这对两（3,3）个阵列有效，并且速度提高了约10倍

我们甚至可以将

nm

重塑为1d，尽管这并不能提高速度：

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N*N)
c3 = np.einsum('n,on,pn->op',A.reshape(N*N),c1,c1)

---

我没有正确解释

opt_einsum

给出的错误

---

问题是，sparse不支持ND sparse数组（它支持！），但我没有使用真正的einsum，因为求和的索引出现了两次以上（

n

和

m

）。如

opt_einsum

文档中所述，这将导致使用sparse.einsum函数，该函数不存在。只使用每个索引中的1或2个就可以了。使用不同的路径，例如hpaulj建议的路径可以用来解决问题。

opt_einsum

是我们很少使用的第三方库（它甚至有

标签吗？）。快速浏览文档并不能表明它是否可以处理
sparse
矩阵
np.einsum
当然不能。@hpaulj
opt_einsum
当然支持
sparse
矩阵，如文档中所示（）：“sparse库也符合要求并受支持…”
c2 = np.einsum('nm,onm,pnm->op',A,c1,c1)

c1 = np.einsum('in,jm->ijnm',B,B).reshape(N*N,N*N)
c3 = np.einsum('n,on,pn->op',A.reshape(N*N),c1,c1)