如何重写JAX代码避免TypeError: unhashable type: 'DynamicJaxprTracer'？

问题背景 #

需要将Python代码改写为JAX代码并通过@jit加速，原有逻辑是用字典将整数映射到函数列表，主函数根据传入的整数索引获取对应函数列表并执行。但添加@jit装饰器后触发错误：

TypeError: unhashable type: 'DynamicJaxprTracer'

错误原因 #

JIT编译时，传入的index参数会被转换为DynamicJaxprTracer类型（JAX用于追踪动态值的对象），而Python字典要求键是可哈希的静态值，因此无法用动态追踪器作为字典键进行索引。

解决方案：使用JAX原生控制流操作 #

JAX提供了jax.lax.switch、jax.lax.cond等原生控制流API，这些操作能被JIT编译器正确处理，替代字典索引实现动态分支选择。

方案1：针对少量键用jax.lax.cond #

如果字典的键数量较少（比如2个），可以直接用cond实现分支判断：

from jax import jit, lax

@jit
def evaluate_functions(xval, index):
    return lax.cond(
        index == 1122997037,
        lambda x: (x**2, 2*x),
        lambda x: (x**3, 3*x),
        xval
    )

print(evaluate_functions(2, 1122997037))  # 输出 (4, 4)
print(evaluate_functions(2, 1124279607))  # 输出 (8, 6)

方案2：针对多键用jax.lax.switch #

如果字典键数量较多，建议用switch实现分支选择，步骤如下：

1. 将字典的键和对应的函数分支整理为列表
2. 在JIT函数中动态判断输入index对应的分支索引
3. 用switch执行对应分支的函数逻辑

示例代码：

from jax import jit, lax

# 原始字典
test_dict = {1122997037: [lambda x: x**2, lambda x: 2*x],
             1124279607: [lambda x: x**3, lambda x: 3*x]}

# 整理键列表和分支函数列表
keys = list(test_dict.keys())
# 每个分支函数接收xval，返回两个函数的执行结果
branches = [lambda x: (f1(x), f2(x)) for f1, f2 in test_dict.values()]

@jit
def evaluate_functions(xval, index):
    # 动态查找index对应的分支索引
    idx = 0
    # 依次判断匹配的键，更新索引
    idx = lax.cond(index == keys[1], lambda: 1, lambda: idx)
    # 若有更多键，继续添加对应的lax.cond判断
    
    # 通过switch选择对应分支执行
    return lax.switch(idx, branches, xval)

print(evaluate_functions(2, 1122997037))  # (4, 4)
print(evaluate_functions(2, 1124279607))  # (8, 6)

方案3：多键场景下用scan自动匹配索引 #

如果键的数量很多，手动写cond判断会很繁琐，可以用jax.lax.scan遍历键列表自动匹配索引：

from jax import jit, lax

test_dict = {1122997037: [lambda x: x**2, lambda x: 2*x],
             1124279607: [lambda x: x**3, lambda x: 3*x]}

keys = list(test_dict.keys())
branches = [lambda x: (f1(x), f2(x)) for f1, f2 in test_dict.values()]

# 定义scan的迭代函数：检查当前键是否匹配目标index
def match_key(carry, key):
    current_idx, target_idx = carry
    # 匹配成功则更新索引为current_idx，否则保持原索引
    matched_idx = lax.cond(target_idx == key, lambda: current_idx, lambda: matched_idx)
    return (current_idx + 1, target_idx), matched_idx

@jit
def evaluate_functions(xval, index):
    # 遍历键列表，找到匹配的分支索引
    (_, _), idx = lax.scan(match_key, (0, index), keys)
    return lax.switch(idx, branches, xval)

print(evaluate_functions(2, 1122997037))  # (4, 4)
print(evaluate_functions(2, 1124279607))  # (8, 6)

核心思路 #

JIT编译的函数中禁止使用动态值作为字典键（因为动态值会被转换为不可哈希的追踪器），必须用JAX原生的控制流操作替代Python的字典索引、if/else等分支逻辑，才能让JIT编译器正确优化代码。

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