python之将 Numpy 中的矩阵列表相乘

我正在寻找一种在 Numpy 中乘以矩阵列表的有效方法。我有一个这样的矩阵:

import numpy as np 
a = np.random.randn(1000, 4, 4) 

我想沿长轴进行矩阵乘法，因此结果是一个 4x4 矩阵。很明显我可以做到:

res = np.identity(4) 
for ai in a: 
    res = np.matmul(res, ai) 

但这太慢了。有没有更快的方法(也许使用 einsum 或其他一些我还不完全理解的函数)？

请您参考如下方法:

需要 log_2(n) for 大小为 2 的幂的堆栈的循环交互的解决方案可能是

while len(a) > 1: 
    a = np.matmul(a[::2, ...], a[1::2, ...]) 

本质上是迭代地将两个相邻矩阵相乘，直到只剩下一个矩阵，每次迭代执行剩余乘法的一半。

res = A * B * C * D * ...         # 1024 remaining multiplications 

成为

res = (A * B) * (C * D) * ...     # 512 remaining multiplications 

成为

res = ((A * B) * (C * D)) * ...   # 256 remaining multiplications 

等等

对于非 2 的幂，您可以对第一个 2^n 矩阵执行此操作，并对其余矩阵使用您的算法。