为何JAX即时编译（JIT）函数比非JIT版本慢很多？

问题场景 #

将纯Python代码改写为JAX代码时，发现一个需要频繁调用的简单函数，其JIT编译版本运行速度远慢于非JIT版本：

import jax.numpy as jnp
from jax import jit

def regular(M,R,a):
    return (3+a)*M*R**a / (4*jnp.pi * R**(3+a))

@jit
def jitted(M,R,a):
    return (3+a)*M*R**a / (4*jnp.pi * R**(3+a))

%timeit regular(1e10,100.,-2.)
# 每次循环346纳秒 ± 2.07纳秒（7次运行的均值±标准差，每次1,000,000次循环）

%timeit jitted(1e10,100.,-2.)
# 每次循环4.2微秒 ± 10.6纳秒（7次运行的均值±标准差，每次100,000次循环）

核心原因 #

你的JIT版本跑得慢，本质是这个函数的运算场景刚好不适合JIT优化：

• JIT调度开销远超计算本身：JIT编译后的函数每次调用，都要经过JAX的调度流程——检查输入形状/类型是否匹配缓存的编译结果、将任务交给XLA执行、返回结果。这些步骤的固定开销（几微秒级别），对于仅需几百纳秒就能完成的标量计算来说，直接把整体耗时拉高了一个数量级。
• 标量运算无法发挥XLA优势：JAX的JIT+XLA是为大规模数组运算设计的，能通过向量化、运算融合等优化大幅提升效率。但你的函数只是几个标量的乘除幂运算，XLA没有优化空间，反而要为JIT的规范付出额外成本。而普通版本的regular调用的是JAX针对标量优化的底层实现，没有JIT的额外负担，自然更快。

解决建议 #

• 如果函数始终处理标量输入，没必要用JIT，直接保留非JIT版本即可。
• 如果后续会切换为批量数组输入（比如M、R、a都是数组），JIT版本的速度会反超非JIT版本——可以尝试将输入换成数组测试，比如regular(jnp.array([1e10]), jnp.array([100.]), jnp.array([-2.]))，就能看到明显的性能差距反转。

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