如何JIT含掩码数组的代码并避免NonConcreteBooleanIndexError？

这是《Count onto 2D JAX coordinates of another 2D array》的后续问题。现有示意函数如下，功能是从卷积结果中提取掩码指定的坐标值，再还原为原尺寸数组：

# 非JIT环境下可正常运行
def predict(model, x1, x2, x2_mask):
    y = somefunc(x1)  # somefunc是卷积操作，x1为棋盘网格编码
    z = y.at[x2[x2_mask][:, 0], x2[x2_mask][:, 1]].get()  # 提取掩码选中的坐标对应值
    w = something.at[jnp.where(x2_mask)].set(z)  # 将z还原为原尺寸N×2数组，其余位置为占位值
    return w

其中：

• x2形状为(N, 2)，x2_mask形状为(N,)
• z的长度由x2_mask中True的数量决定，w固定为(N,2)

报错场景 #

1. nnx.jit编译报错：添加@nnx.jit装饰器后触发NonConcreteBooleanIndexError，移除掩码逻辑后JIT可正常运行。
2. vmap批处理报错：使用jax.vmap封装后，赋值语句触发ConcretizationTypeError。

核心矛盾：输入尺寸固定，y和w的大小可推断，仅z长度随掩码变化，但JIT编译仍因动态形状报错。

JAX的JIT要求所有数组形状在编译时确定，原代码中x2[x2_mask]会产生形状依赖掩码的数组，导致编译时无法确定形状约束，进而触发错误。以下是两种可行的改造思路：

方案1：用广播+掩码过滤替代动态索引 #

直接提取所有x2坐标对应的y值，再通过掩码筛选保留目标值，全程使用固定形状数组：

@nnx.jit
def predict(model, x1, x2, x2_mask):
    y = somefunc(x1)
    # 提取所有x2坐标对应的y值，形状固定为(N, 2)
    all_vals = y[x2[:, 0], x2[:, 1]]
    # 用掩码选择保留目标值或原始占位值
    w = jnp.where(x2_mask[:, None], all_vals, something)
    return w

优势 #

• 完全规避动态形状数组，完美适配JIT编译要求
• 代码简洁，原生兼容jax.vmap，无需额外修改

方案2：用jax.lax动态操作处理可变长度场景 #

如果必须保留可变长度的z数组（比如后续有依赖z长度的逻辑），可通过jax.lax的显式动态操作约束形状：

@nnx.jit
def predict(model, x1, x2, x2_mask):
    y = somefunc(x1)
    # 获取固定长度的掩码索引数组（未选中位置填充0，不影响赋值）
    indices = jnp.where(x2_mask, size=x2_mask.shape[0])[0]
    # 动态提取掩码选中的x2坐标
    selected_x2 = jax.lax.dynamic_slice_in_dim(x2, 0, jnp.sum(x2_mask), axis=0)
    # 提取对应y值
    z = y[selected_x2[:, 0], selected_x2[:, 1]]
    # 动态填充到目标数组
    w = something.at[indices].set(z)
    return w

原理 #

• jnp.where(..., size=N)强制返回固定长度的索引数组，消除JIT的形状不确定性
• jax.lax.dynamic_slice_in_dim显式指定切片长度，让编译器能推断操作的形状约束

vmap批处理适配 #

上述两种方案均直接兼容jax.vmap，批处理代码可直接使用：

batched_predict = jax.vmap(predict, in_axes=(None, 0, 0, 0))
batched_predict(model, x1s, x2s, x2_masks)

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