np.linalg.inv不支持批量矩阵求逆,因其仅接受2D输入,对(N,M,M)数组会报“Expected 2D array, got 3D array”错误;应改用np.linalg.solve逐个求解或借助vectorize、torch.inverse等支持广播的方案。

为什么不能直接用 np.linalg.inv 处理批量矩阵
当你有一组形状为 (N, M, M) 的矩阵(比如 1000 个 3×3 矩阵),直接对整个数组调用 np.linalg.inv 会报错:LinAlgError: Expected 2D array, got 3D array instead。NumPy 的 np.linalg.inv 只接受二维输入,不支持广播式批量求逆。
用 np.linalg.solve 替代:更稳定也更快
批量求逆的本质是解 A_i @ X_i = I,即对每个 A_i 求其逆矩阵 X_i。用 np.linalg.solve 显式求解比先算逆再乘更数值稳定,也避免构造单位矩阵的内存开销。
实操建议:
- 对每个
i,调用np.linalg.solve(A[i], np.eye(M))—— 简单但慢(Python 循环) - 改用
np.vectorize包裹np.linalg.solve并指定signature='(m,m),(m)->(m,m)',可启用 NumPy 的隐式广播(需 NumPy ≥ 1.20) - 更推荐:用
scipy.linalg.block_diag拼成大块对角矩阵再整体求逆(仅适用于小M、中等N) - 若
M固定且较小(如 2–4),手写解析公式(如 2×2 的(1/det) * [[d, -b], [-c, a]])反而最快,无依赖、无分支、全向量化
使用 torch.inverse 或 jax.numpy.linalg.inv 的前提条件
如果你的场景允许引入额外依赖,PyTorch 和 JAX 原生支持 (..., M, M) 批量求逆:
立即学习“Python免费学习笔记(深入)”;
-
torch.inverse(A)要求A是torch.Tensor且在 GPU 上时加速明显;但需注意:输入必须是非奇异的,否则抛出RuntimeError -
jax.numpy.linalg.inv(A)支持 JIT 编译和自动微分,但对A的 batch size 过大时可能触发 JIT 编译超时或显存爆炸 - 二者都不检查条件数,坏矩阵会导致结果不可靠,而 NumPy 的
np.linalg.inv至少会在单矩阵上抛出LinAlgError
自己实现批处理时最容易忽略的三个点
手动循环 + np.linalg.inv 看似简单,但实际部署常翻车:
- 忘记预检查奇异矩阵:用
np.linalg.det(A[i]) == 0不可靠(浮点误差),应改用np.linalg.cond(A[i]) > 1e15或捕获LinAlgError - 输出数组未预分配:边循环边
list.append()再np.stack(),会触发多次内存拷贝;应提前inv_A = np.empty_like(A) - 没设 dtype 一致性:输入是
float32却用默认float64求逆,不仅慢还可能溢出;务必传dtype=A.dtype给np.eye和求解器
批量矩阵求逆没有银弹——小规模手写公式,中等规模用 np.vectorize + solve,大规模且可换生态就上 JAX/PyTorch;但无论哪种,条件数监控和 dtype 对齐几乎总是被跳过的一步。


















