JAX中如何获取JIT函数重编译通知，有没有可靠实现方案？

JAX官方已经提供了稳定的重编译监听和拦截能力，不需要依赖tracer执行副作用的内部hack。

1 重编译通知的稳定实现方案 #

1.1 本地调试用：开启编译日志 #

如果只是调试阶段需要感知重编译触发，直接开启官方的编译日志配置即可，不需要修改业务代码：

import jax
# 全局开启重编译日志
jax.config.update("jax_log_compiles", True)

也可以在启动脚本前设置环境变量 export JAX_LOG_COMPILES=1，所有JIT函数重编译时都会打印触发原因、输入特征、编译耗时等详细信息。

1.2 生产/二次开发用：注册自定义编译回调 #

从JAX 0.4.6版本开始，官方提供了公开的编译回调注册接口，你可以绑定任意自定义逻辑，重编译触发时自动执行，完全不依赖内部实现逻辑：

import jax
from jax import CompilationArtifact

recompilation_count = 0
# 自定义允许的最大重编译次数，超过直接拦截
MAX_RECOMPILATION_LIMIT = 3

def compilation_callback(artifact: CompilationArtifact):
    global recompilation_count
    recompilation_count += 1
    # 可在这里加日志、监控上报逻辑
    print(f"重编译触发：函数名={artifact.function_name}，输入特征={artifact.input_signature}，当前累计重编译次数={recompilation_count}")
    # 超过阈值直接抛出异常拦截编译
    if recompilation_count > MAX_RECOMPILATION_LIMIT:
        raise RuntimeError(f"函数{artifact.function_name}重编译次数超过阈值{MAX_RECOMPILATION_LIMIT}，已终止编译，请检查输入是否存在频繁变化的动态维度/静态参数")

# 全局注册回调，所有JIT编译都会触发
jax.register_compilation_callback(compilation_callback)

# 测试代码
@jax.jit
def func(z):
    return z * z + 100 / z

func(1)
func(2)
func(jax.numpy.arange(10))

2 频繁重编译的预防方案 #

大部分频繁重编译都是因为没有正确标记静态参数、或者输入形状频繁变化导致的，可以针对性处理：

• 显式标记静态参数：如果JIT函数的部分参数不会频繁变化，用 @jax.jit(static_argnums=(0,), static_argnames=("config",)) 标记这些参数，只有参数值真的发生变更时才会触发重编译
• 固定/标注动态形状：如果输入形状存在动态变化的维度，可以用 jax.jit(dynamic_shapes=True) 或者显式指定可变维度的范围，避免每一种新形状都触发一次重编译
• 限制缓存大小：设置 jax.config.update("jax_max_trace_cache_size", 8)，限制单个JIT函数最多缓存多少个编译版本，超过阈值后自动淘汰最早的编译结果，避免内存占用过高

你之前使用的全局变量计数的方案确实不推荐长期使用，一方面JAX未来版本可能会调整tracer阶段的副作用处理规则，另一方面多线程/分布式场景下全局变量计数会出现偏差，且无法获取重编译的具体原因做针对性优化。

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