如何在Python中无需循环实现N个方阵的类似numpy.matmul运算？

numpy 本身就提供了无需循环的批量矩阵乘法方案，分两种常见场景处理：

1. 对应位置的方阵逐对相乘 #

如果你有两组N个方阵（比如数组形状为 (N, M, M)，每个元素是一个M×M方阵），直接使用 @ 运算符或 np.matmul 就能完成批量逐对相乘，完全不需要循环。

示例代码：

import numpy as np

# 生成3个2×2的方阵，存储为形状(3, 2, 2)的数组
batch_A = np.random.rand(3, 2, 2)
batch_B = np.random.rand(3, 2, 2)

# 批量计算每个对应位置的矩阵乘积，结果形状为(3, 2, 2)
batch_result = batch_A @ batch_B
# 等价于 np.matmul(batch_A, batch_B)

2. N个方阵按顺序连乘得到单一矩阵 #

如果需要将N个方阵依次相乘（如 A1 × A2 × ... × AN），可以使用 np.linalg.multi_dot，它支持传入矩阵列表直接计算连乘，内部做了优化，比手动循环更高效。

示例代码：

# 将N个方阵存入列表（这里用batch_A中的3个方阵为例）
matrix_list = [batch_A[0], batch_A[1], batch_A[2]]

# 计算连乘结果，得到一个2×2矩阵
final_matrix = np.linalg.multi_dot(matrix_list)
# 等价于 batch_A[0] @ batch_A[1] @ batch_A[2]

补充说明 #

• 对于更高维度的批量数据（比如形状为 (K, N, M, M)，代表K组每组N个M×M方阵），np.matmul 同样能处理，只要保证维度对齐即可。
• 这些方法都是numpy内置的向量化操作，性能远高于手动循环，尤其当N较大时优势明显。

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