JAX提取二维数组行内非NaN值时遭遇NonConcreteBooleanIndexError问题

给定含NaN的JAX二维数组，需每行保留非NaN值（保持原顺序），结果行数不变，列数设为每行最多非NaN值数量k，不足部分用NaN填充。直接用vmap处理逐行函数会因动态尺寸触发错误，以下是高效解决方案：

错误原因 #

JAX基于静态图框架，vmap要求映射函数内部的操作形状在编译时确定。原代码中row_vals[row_mask]返回的数组长度随每行非NaN数量动态变化，无法被静态编译器处理，因此触发NonConcreteBooleanIndexError。

高效解决方案 #

采用向量化的掩码累积和方法，全程保持静态形状，避免动态尺寸操作，适合大规模数组处理：

import jax.numpy as jnp

array_2d = jnp.array([
    [jnp.nan,        1,       2,   jnp.nan,    3],
    [10     ,jnp.nan,   jnp.nan,        20,jnp.nan]
])

k = 3  # 每行最多保留的非NaN值数量

# 1. 生成非NaN值的掩码
mask = ~jnp.isnan(array_2d)

# 2. 计算每行非NaN值的累积计数，确定每个值在输出中的目标列位置
target_cols = jnp.cumsum(mask, axis=1) - 1  # 转换为0起始索引

# 3. 筛选出需要保留的前k个非NaN值的位置
valid_pos = mask & (target_cols < k)

# 4. 初始化输出数组为NaN，将有效位置的值填入对应列
output = jnp.full((array_2d.shape[0], k), jnp.nan)
output = output.at[jnp.where(valid_pos)[0], target_cols[valid_pos]].set(array_2d[valid_pos])

print(output)

输出结果 #

[[ 1.  2.  3.]
 [10. 20. nan]]

方案优势 #

• 全程向量化操作，无动态分支或循环，性能高效，适配大规模数组（即使k=10或100也能快速运行）
• 完全符合JAX静态图要求，可安全用于vmap、jit等转换
• 严格保持原数组中非NaN值的顺序

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