添加整数参数至损失函数后JAX/Equinox训练流水线性能骤降

我基于JAX和Equinox搭建了训练流水线，希望将批次索引传入损失函数，以根据索引执行不同逻辑。未传入批次索引时训练循环耗时约15秒，但传入索引后耗时骤增至约1小时。我是JAX新手，想知道问题出在哪。

我的训练流水线代码：

def fit_cv(model: eqx.Module, 
           dataloader: jdl.DataLoader, 
           optimizer: optax.GradientTransformation, 
           loss: tp.Callable, 
           n_steps: int = 1000):
    
    opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
    dloss = eqx.filter_jit(eqx.filter_value_and_grad(loss))
    
    @eqx.filter_jit
    def step(model, data, opt_state, batch_index):
        loss_score, grads = dloss(model, data, batch_index)
        updates, opt_state = optimizer.update(grads, opt_state)
        model = eqx.apply_updates(model, updates)
        return model, opt_state, loss_score
    
    loss_history = []
    for batch_index, batch in tqdm(zip(range(n_steps), dataloader), total=n_steps):
        if batch_index >= n_steps:
            break
        batch = batch[0] # dataloader returns tuple of size (1,)
        model, opt_state, loss_score = step(model, batch, opt_state, batch_index)
        loss_history.append(loss_score)
    return model, loss_history

损失函数签名：

def loss(self, model: eqx.Module, data: jnp.ndarray, batch_index: int):

我的需求是在N步后切换两种损失函数，因此需要获取批次索引的具体值。

原因分析 #

问题出在JIT编译的特性上：你当前传入step函数的batch_index是Python整数，而JAX的JIT会将Python标量视为编译常量。这意味着每一个不同的batch_index都会触发一次全新的JIT编译流程——1000步训练就会编译1000次，每次编译都需要耗费大量时间，最终导致总训练时间暴增。

解决方案 #

要解决这个问题，需要做两点：

1. 将batch_index转换为JAX数组（比如jnp.int32(batch_index)）传入JIT函数，让它成为可追踪的张量而非编译常量；
2. 使用jax.lax.cond替代Python原生的条件判断，确保逻辑能被JAX正确追踪和编译，避免分支导致的重复编译。

修改后的损失函数核心逻辑示例：

condition = (batch_index // self.switch_steps) % 2 == 1
# 确保condition是JAX数组类型
condition = jnp.asarray(condition)
loss_value = jax.lax.cond(
    condition,
    lambda: loss1(inputs),
    lambda: loss2(inputs),
)
return loss_value

同时，在训练循环中传入索引时要转换为JAX数组：

# 原代码中的调用行修改为
model, opt_state, loss_score = step(model, batch, opt_state, jnp.int32(batch_index))

这样JIT只会编译一次step函数，后续所有批次都复用这个编译结果，训练耗时就会回到正常水平。

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