PAYATHON 2026

JAX 如何处理扩展窗口以避免按行数重复编译？

数据小许

结论

要避免扩展窗口每增加一行就触发一次编译，传给 jax.jit 的数组形状必须保持不变。可以将完整数据（或预设上限大小的数据）作为固定形状输入，再用布尔掩码表示当前时间窗口。模型计算时通过掩码排除窗口外的数据。

原示例中的 X[X[:, 0] <= t, ...] 会根据数组中的值筛出不定行数的结果。JAX 在 jit 跟踪期间无法确定结果形状，因此会触发 NonConcreteBooleanIndexError。即使在 JIT 函数外筛选，只要 X_filtered 的行数发生变化，输入形状也会随之变化，需要重新编译。

固定形状输入与掩码

逐行计算时，可以在 JIT 函数中保留所有行，并用掩码标记当前窗口内有效的行：

import jax
import jax.numpy as jnp

@jax.jit
def test_func(X, t):
    in_window = X[:, 0] <= t
    values = X[:, 1:] ** 2
    return jnp.where(in_window[:, None], values, 0)

X = jnp.array([
    [0, 1, 1],
    [0, 1, 1],
    [1, 2, 4],
    [2, 3, 9],
    [2, 3, 9],
    [3, 4, 16],
    [4, 5, 25],
])

for t in range(5):
    result = test_func(X, jnp.asarray(t))
    print(result)

每次调用传入的 X 形状都相同，t 是一个标量数组。结果会保留所有输入行，并将窗口外的行置为零。如果零本身也可能是有效数据，就要同时保留并使用 in_window 掩码，不能只根据输出值判断行是否有效。

将掩码用于模型目标

如果模型通过逐行损失来估计参数，可以先计算每行的损失，再用掩码汇总：

@jax.jit
def masked_mean_loss(prediction, target, in_window):
    row_loss = (prediction - target) ** 2
    row_loss = jnp.where(in_window, row_loss, 0.0)
    count = jnp.maximum(in_window.sum(), 1)
    return row_loss.sum() / count

实际使用时，应将掩码传入目标函数、似然或统计量计算，确保窗口外的观测不会影响参数估计。masked_mean_loss 展示的是一种掩码汇总方式，不能直接套用到所有统计模型中。

如果模型必须接收紧凑的变长矩阵，或包含依赖当前样本数的运算，就可能需要重新设计计算方式。可以将行数填充到有限个固定大小的 bucket，并在模型内屏蔽填充行。这样，编译次数最多与 bucket 数量相关。如果固定形状和分桶都不可行，就需要承担不同输入形状分别编译的成本。

关于重编译

JAX 会按输入类型和形状缓存编译结果。输入形状不同，通常需要重新跟踪和编译。将完整数据保持为同一形状，并把时间索引作为数组值传入，有助于复用同一编译结果。也要避免在循环中反复创建新的 JIT 函数；将参数声明为静态参数时，不同取值也可能触发重新编译。JAX JIT 文档

扩展窗口应通过掩码值的变化来表示，传入 JIT 函数的数组行数则保持不变。JAX 限制动态布尔索引，是因为编译时必须确定数组形状；掩码运算可以保持形状静态。JAX 动态形状说明；NonConcreteBooleanIndexError 说明

备注：内容仅供参考。