运行时重定义JAX-JIT函数的子函数后，如何最优重编译主函数？

问题场景 #

原始代码如下：

from jax import jit

def bar(x):
  return x ** 2

@jit
def foo(x):
  return 1 + bar(x)

print(f'foo(4) = {foo(4)}')  # 输出 foo(4) = 17

当运行时重定义bar后：

def bar(x):
  return 2 * x

print(f'foo(4) = {foo(4)}')  # 仍输出 foo(4) = 17，因为foo使用的是旧bar的编译结果

重新编译foo的最优方式 #

方案1：提前保留未jit的原始函数（推荐） #

定义时将业务逻辑与jit装饰分离，保留原始未编译的函数引用：

from jax import jit

def bar(x):
  return x ** 2

# 保留未jit的核心逻辑函数
def _foo(x):
  return 1 + bar(x)

# 基于原始函数创建jit版本
foo = jit(_foo)

当bar更新后，只需重新执行jit包装即可，无需重写函数体：

def bar(x):
  return 2 * x

# 重新编译原始函数，使用最新的bar
foo = jit(_foo)
print(f'foo(4) = {foo(4)}')  # 输出 foo(4) = 9

方案2：从已jit函数中恢复原始函数 #

如果没有提前保留原始函数，可以通过__wrapped__属性获取被jit装饰的原始函数，再重新编译：

def bar(x):
  return 2 * x

# 从已jit的foo中提取原始函数，重新编译
foo = jit(foo.__wrapped__)
print(f'foo(4) = {foo(4)}')  # 输出 foo(4) = 9

不推荐的方案 #

直接执行foo = jit(foo)虽然在当前场景能运行，但本质是对已jit的函数再次进行jit编译，相当于编译“调用jit版foo”的逻辑，而非直接基于最新的bar重新编译原始业务逻辑，存在潜在的异常风险（比如多层jit包装导致的性能损耗或编译逻辑偏差），因此不建议使用。

额外问题：能否自动重新编译所有依赖bar的函数？ #

不能。JAX的jit编译是静态快照式的：当你jit编译foo时，JAX会将当时bar的函数体内联到foo的编译结果中，不会在运行时追踪Python函数引用的动态变化。由于Python的函数是动态对象，JAX无法自动检测到bar的重定义，也无法自动触发所有依赖它的jit函数重新编译。

在复杂项目中，建议通过集中管理核心业务函数+统一jit包装入口的方式，来避免遗漏需要重新编译的函数。

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