JAX/JIT与标准NumPy性能接近：辛普森积分场景问题出在哪里？

结论 #

你的NumPy代码本身效率极高，JAX CPU端性能接近属于正常现象，你的JAX代码写法没有问题。两个版本计算结果的最大相对误差仅8×10^-16，也进一步验证了JAX实现的正确性。

具体原因 #

• 你的纯NumPy实现全程采用向量化操作，linspace、切片赋值、einsum等接口都直接调用了底层高度优化的BLAS/LAPACK库，本身就是接近C语言的运行效率，CPU上已经很难有进一步的性能提升空间。
• JAX的CPU后端基于XLA编译，对于这种已经高度向量化的数值计算，XLA能做的优化空间非常有限，再加上JAX本身的少量调度开销，最终性能和NumPy拉不开差距是非常正常的。
• GPU上的巨大性能差距来自于并行计算优势：你当前的计算任务是200个独立被积函数、100组独立积分区间，属于天然高度并行的计算负载，NumPy只能运行在CPU上无法利用GPU并行能力，而JAX可以将全部运算卸载到GPU执行，因此性能提升超过100倍，这也反过来验证了你的JAX实现逻辑正确。

CPU端JAX性能优化建议 #

如果想要进一步提升JAX在CPU上的运行速度，可以尝试以下调整：

1. 开启XLA快速数学优化，运行代码前设置环境变量：

export XLA_FLAGS=--xla_cpu_enable_fast_math

2. 优化jax_funcN的实现，避免Python侧的循环编译展开，改用广播向量化写法：

@jit
def jax_funcN(x):
    i = jnp.arange(200)
    return (x[..., None] ** (i / 10)) * jnp.exp(-x[..., None])

3. 若使用的是ARM架构CPU，可以尝试开启XLA的SVE指令集优化，能进一步提升向量化计算效率。

内容的提问来源于stack exchange，提问作者Jean-Eric