JAX Jit类方法装饰器与直接调用的差异原因探究

熟悉JAX的开发者大多见过官方文档里的类方法jit示例，也知道下面这段代码无法运行：

import jax.numpy as jnp
from jax import jit

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  @jit  # <---- 这里怎么正确实现？
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

c = CustomClass(2, True)
c.calc(3)  

官方文档给出了3种解决方案，但实际测试发现，直接把jit作为函数调用而非装饰器使用时，代码能正常运行，JAX不会报错无法处理CustomClass类型的self：

import jax.numpy as jnp
from jax import jit

class CustomClass:
  def __init__(self, x: jnp.ndarray, mul: bool):
    self.x = x
    self.mul = mul

  # 此处无装饰器！
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

c = CustomClass(2, True)
jitted_calc = jit(c.calc)
print(jitted_calc(3))

运行输出：

6 # 运行正常！

进一步测试发现，这种方式的行为等价于用@partial(jax.jit, static_argnums=0)把self标记为静态参数——修改self的属性不会影响后续调用结果：

c = CustomClass(2, True)
jitted_calc = jit(c.calc)
print(jitted_calc(3))
c.mul = False 
print(jitted_calc(3))

运行输出：

6
6 # 结果无更新

对比普通装饰器的行为，普通装饰器会正常响应self属性的修改：

def decorator(func):
    def wrapper(*args, **kwargs):
        x = func(*args, **kwargs)
        return x
    return wrapper

custom = CustomClass(2, True)
decorated_calc = decorator(custom.calc)
print(decorated_calc(3))
custom.mul = False
print(decorated_calc(3))

运行输出：

6
3

问题 #

为何JAX的jit在两种使用方式下表现差异极大？直接调用jit(c.calc)时为什么能处理self类型？

原因解析 #

这两种使用方式的核心差异在于**self参数的绑定时机与JAX对参数的处理逻辑**：

• 装饰器方式的问题
  用@jit装饰类方法时，装饰器是在类定义阶段生效的，此时calc还只是一个未绑定的方法（属于类的属性），没有和具体实例self绑定。JAX的jit会尝试追踪这个未绑定方法的参数类型，而未绑定方法的第一个参数self是CustomClass类型，不属于JAX能处理的可追踪数组或静态参数范畴，因此会抛出类型错误。

• 直接调用jit(c.calc)的逻辑
  当你执行jit(c.calc)时，c.calc已经是一个绑定方法——Python自动把实例c作为self参数绑定到了方法上，此时传给jit的函数实际上是一个只接受y作为参数的“偏函数”：它的self已经固定为实例c，不再是函数的显式参数。
  JAX处理这种绑定方法时，会把绑定的self视为静态参数，在编译时就捕获self当时的所有属性值（比如self.x=2、self.mul=True），并把这些值硬编码到编译后的XLA计算图中。后续修改self的属性时，已经编译好的计算图不会更新，所以调用结果不会变化——这和用static_argnums=0标记未绑定方法的self参数效果完全一致。

• 普通装饰器的差异
  普通装饰器只是对传入的函数（这里是绑定方法c.calc）做了一层包装，每次调用装饰后的函数时，都会实际执行原绑定方法，而原绑定方法会实时读取当前实例self的属性值，所以修改self属性后结果会更新。JAX的jit则是直接编译了绑定方法当时的计算逻辑，后续调用不会再读取最新的self属性。

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