如何高效对列表中可变数量的矩阵进行连续相乘？

解决方案 #

np.linalg.multi_dot 本身支持直接传入包含任意数量矩阵的列表对象，不需要手动枚举每个元素的索引，直接将存储矩阵的列表作为参数传入即可。

代码示例 #

import numpy as np

# 代码开头可自由设置矩阵数量
matrix_count = 6
# 示例：生成包含matrix_count个矩阵的列表a
a = [np.random.rand(3, 3) for _ in range(matrix_count)]

# 直接传入列表完成连续相乘，无需手动写每个索引
M.append(np.linalg.multi_dot(a))

其他可选实现 #

如果需要自定义相乘逻辑，也可以用functools.reduce实现，不过该方式不会自动优化相乘顺序，大尺寸矩阵场景下效率低于multi_dot：

from functools import reduce
# 两种写法等价
M.append(reduce(np.dot, a))
M.append(reduce(lambda x, y: x @ y, a))

内容的提问来源于stack exchange，提问作者Nrmn