结合 sparse 和 einsum 进行大的稀疏求和

Combining sparse and einsum to perform large sparse sum

我有一个形状为 (N, N) 的矩阵 A 和一个形状为 (N, N) 的矩阵 B。 我正在使用以下 einsum 构建矩阵 M(使用 opt_einsum 库):

M = oe.contract('nm,in,jm,pn,qm->ijpq', A, B, B, B, B)

这是计算以下总和:

这产生了形状为 (N, N, N, N) 的矩阵 M。然后我将其重塑为形状为 (N**2, N**2)

的二维数组
M = M.reshape((N**2, N**2))

这必须是二维的,因为它被视为线性运算符。

我想使用 sparse 库,因为 M 是稀疏的,并且变得太大而无法存储大 N。我可以使 A 和 B 稀疏,并将它们插入 oe.contract .

问题是,稀疏仅支持二维数组,因此无法生成形状为 (N, N, N. N) 的 4D 输出。有没有办法结合 einsum 和 reshape 步骤以允许以这种方式使用稀疏,因为 M 的最终形状是 2D?

这可能对您使用 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给出的错误。

问题是不是稀疏不支持 ND 稀疏数组(它支持!),但我没有使用真正的 einsum,因为求和的索引看起来更多超过两次(nm)。如 opt_einsum 文档中所述,这将导致使用 sparse.einsum 函数,其中 none 存在。每个索引仅使用 1 或 2 个即可。使用不同的路径,例如 hpaulj 建议的路径可以用来解决问题。