以对角矩阵作为向量存储的多个数组

Question

我有一个一维数组A= a, b, c...和三维数组T的形状(N_A，N_B，N_A)。A是指用N_A矩阵表示对角线N_A。

我想用T来做A的收缩，而不需要把A提升到高密度的存储。特别是，我想

np.einsum('ij, ikl', A, T)

和

np.einsum('ikl, lm', T, A)

在保持稀疏的情况下，能做这样的事情吗？

注意:这个问题类似于

dot product with diagonal matrix, without creating it full matrix

但不完全相同，因为我不清楚如何推广到更复杂的索引模式。

Answer 1

np.einsum('ij, ikl', np.diag(a), t)等同于(a * t.T).T。

np.einsum('ikl, lm', t, np.diag(a))等同于a * t。

(经反复试验发现)