NumPy 连续矩阵乘法向量化

Question

我有一个形状为

ar

的 NumPy 数组

(M, L, N, N)

。

我想将

L

(N, N)

矩阵依次相乘 (

multiplied_ar[m] = ar[m, 0, :, :] @ ar[m, 1, :, :] @ ...

) 以获得形状为

(M, N, N)

的数组。

是否可以以某种方式对其进行向量化，这样我就不必迭代

m

和

l

？

Answer 1

你可以使用Python的

functools.reduce(func,seq)

和

np.einsum

:

import functools
a=np.random.rand(M,L,N,N)
b = np.transpose(a, (1, 0, 2, 3))
def m_einsum(A1, A2):
  return np.einsum('fij,fjk->fik', A1, A2)

res=functools.reduce(m_einsum, b)