NumPy中批量轴为何前置？切换约定需考量性能问题

问题背景 #

我开发的所有包都将后置轴用作批量轴，这在我看来更符合直觉，但NumPy始终把前置轴作为批量轴。现在我考虑切换到NumPy的约定，只为让NumPy用户更容易理解我的代码，想问问有没有什么建议？

性能对比测试 #

从性能角度来看，切换到前置批量轴可能是个糟糕的选择，以下是测试代码及结果：

import numpy as np

np.random.seed(6512)
a = np.random.rand(50000, 8, 3, 3)

np.random.seed(85742)
b = np.random.rand(50000, 8, 3, 3)

c = a @ b
# 每次循环耗时19.8毫秒±543微秒（7次运行，每次100循环的均值±标准差）

d = np.einsum("...ik,...kj->...ij", a, b)
# 每次循环耗时84.1毫秒±2.4毫秒（7次运行，每次10循环的均值±标准差）

# 现在使用后置轴（确保转置后的数据为C连续数组）
A = np.ascontiguousarray(np.transpose(a, [2, 3, 0, 1])) # A_ijab
B = np.ascontiguousarray(np.transpose(b, [2, 3, 0, 1])) # B_ijab

C = (B.T @ A.T).T # (C^T)_baji = B_bajk A_baki -> C_ijab
# 每次循环耗时16.9毫秒±1.82毫秒（7次运行，每次100循环的均值±标准差）

D = np.einsum("ik...,kj...->ij...", A, B)
# 每次循环耗时17.2毫秒±842微秒（7次运行，每次100循环的均值±标准差）

assert np.allclose(c, d)
assert np.allclose(C, D)
assert np.allclose(np.transpose(D, [2, 3, 0, 1]), d)
assert np.allclose(np.transpose(C, [2, 3, 0, 1]), c)

更复杂的einsum运算测试 #

# 交叉并矢积
# ----------------------

E = np.einsum("ik...,jl...->ijkl...", A, B)
# 每次循环耗时76.5毫秒±2.22毫秒（7次运行，每次10循环的均值±标准差）

e = np.einsum("...ik,...jl->...ijkl", a, b)
# 每次循环耗时207毫秒±3.29毫秒（7次运行，每次10循环的均值±标准差）

assert np.allclose(np.transpose(E, [4, 5, 0, 1, 2, 3]), e)

核心建议 #

• 优先权衡性能与兼容性：如果你的包性能敏感（比如涉及大量矩阵运算、einsum操作），从测试结果看后置批量轴的性能优势明显，强行切换会导致运行效率大幅下降，这时候兼容性的代价可能过高。
• 提供转换接口：不需要完全切换约定，而是给用户提供批量轴前后置转换的便捷函数，既保留自己的设计直觉，也让NumPy用户能轻松适配你的包。
• 文档明确约定：在包的文档中清晰标注你使用的批量轴位置，给出示例代码说明如何与NumPy的前置轴风格交互，降低用户的理解成本。
• 考虑生态适配：如果你的包需要深度集成NumPy、PyTorch等主流库（它们大多使用前置批量轴），长期来看统一约定能减少用户的转换代码，但前提是性能损失在可接受范围内。

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