JAX JIT函数：参数与全局变量的差异及实践规则问询

核心区别 #

• 参数传入：JAX会将传入的参数视为动态可追踪的输入。如果参数的形状、dtype不变仅值变化，JAX会复用已编译的缓存；若形状或dtype改变，则触发重新编译。这种方式能正确追踪变量更新，适合动态变化的数值。
• 全局变量引用：JAX在jit编译时会把全局变量当时的取值当作静态常量硬编码进编译后的计算图。编译完成后，全局变量后续的修改不会被jit函数感知，函数会一直使用编译时的旧值。

全局变量修改后的行为 #

举个实际场景：假设你定义了全局变量extra_learning_rate = 0.1，并在jit装饰的step函数里直接使用它。之后你把extra_learning_rate改成0.01再调用step，函数依然会用0.1来计算——因为编译阶段已经把这个值固定死了。除非你手动清除JAX的编译缓存（jax.clear_caches()），或者重新装饰jit函数，否则不会使用新的全局变量值。

结合Optax的实践规则 #

Optax的optimizer涉及两个关键部分，二者处理方式完全不同：

• optimizer定义：如果是固定不变的配置（比如训练全程用Adam优化器、学习率初始值固定），可以作为全局变量，因为它本身是静态的函数/配置对象，不会在训练中动态变化。
• optimizer状态：必须作为参数传入jit函数，并且每次step后更新状态再传入下一次调用。因为状态包含动量、学习率调度的动态信息（比如余弦退火的当前步数），如果把它放在全局变量里，jit编译会缓存初始状态，导致每次step都用初始值，训练完全不会收敛。

高性能开发的关键规则 #

• 动态变量必传参：模型参数、optimizer状态、批次数据、动态调整的超参数（比如随步数变化的学习率），必须作为参数传入jit函数，绝对不能用全局变量。
• 静态变量合理处理：如果某些配置全程不变（比如固定的损失函数、静态开关），可以用全局变量；但如果后续可能需要修改（比如切换优化器），建议作为参数传入，或者用jax.jit的static_argnums/static_argnames标记，这样修改时能触发重新编译，同时避免不必要的编译缓存。
• 减少重复编译：对于静态参数（比如固定的batch size、训练模式开关），用static_argnums标记后，JAX只会在参数的形状/类型变化时重新编译，相同形状的不同值会复用缓存，大幅提升性能。

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