如何无循环计算(K,N,2,2)型ndarray各栈内2x2矩阵乘积

对于维度为(K,N,2,2)的numpy数组，要实现每个K对应的N个2x2矩阵连乘且不使用显式for循环，有几种高效的方法：

方法1：使用np.reduce + np.matmul #

这是最简洁通用的方案，利用np.reduce沿着N所在的轴（axis=1）累积执行矩阵乘法：

import numpy as np

# 生成测试数据
K, N = 3, 4
A = np.random.rand(K, N, 2, 2)

# 计算每个栈内N个矩阵的连乘
result = np.reduce(np.matmul, A, axis=1)

原理说明 #

• np.reduce会遍历axis=1（即每个K对应的N个矩阵），依次对相邻矩阵执行np.matmul操作，最终得到每个K对应的2x2乘积矩阵。
• 这个操作等价于对每个k执行np.linalg.multi_dot(A[k])，但完全避免了显式for循环。

验证正确性 #

可以用for循环的结果对比验证：

# 用for循环计算作为对照
result_loop = np.zeros((K, 2, 2))
for k in range(K):
    result_loop[k] = np.linalg.multi_dot(A[k])

# 检查结果是否一致（浮点误差允许范围内）
print(np.allclose(result, result_loop))  # 输出True

方法2：固定N时用np.einsum（可选） #

如果N的长度是固定的，可以用np.einsum手动写出连乘的维度映射。比如当N=3时：

result_einsum = np.einsum('kabi,kbcj,kcdl->kadi', A[:,0], A[:,1], A[:,2])

但这种方法仅适用于N固定的场景，灵活性不如np.reduce方案。

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