JAX jit函数中如何按条件正确累加jnp.array数组

云上老周

2026-8-27

问题原因

你遇到的问题核心是JAX JIT的纯函数执行规则，以及控制流原语的运行逻辑：

@jit装饰的函数编译时，只会执行一遍Python层面的代码做抽象追踪，用来构建XLA计算图，编译完成后实际在设备上运行的是优化后的HLO代码，根本不会再执行你写的Python赋值语句。
你在body_fun里写的global accu修改，只会在追踪阶段触发一次：也就是JAX第一次跑到循环位置时，把初始形状为(1,3)的accu和第一行待遍历的data[1]拼接，得到形状(2,3)的追踪数组，这个值只存在于追踪上下文里，不会同步到函数外的全局变量。
lax.while_loop的所有循环迭代都在编译后的计算图内部执行，迭代间的状态必须通过carry值显式传递，不支持向外透传Python侧的副作用。你看到的那个带DynamicJaxprTrace标记的对象，就是追踪阶段生成的临时抽象值，等JIT函数执行完就会失效，根本不是实际运算得到的结果数组。
你的循环逻辑本身是对的：初始索引是1，判断data[1]全非零就进循环，索引涨到2，再判断data[2]全非零进循环，索引涨到3时data[3]全零，循环终止，本来应该累加索引1、2两行，加上初始的第0行总共3行，但因为全局变量副作用不生效，你最后拿到的还是追踪阶段只拼了一次的(2,3)形状的临时值。

最优实现方案

完全不需要用全局变量，JAX里做这类筛选累加有两种标准写法，都符合纯函数要求，没有副作用问题：

方案一：直接用布尔索引筛选（性能最优，优先用）
如果只是筛选全非零的行，不需要写显式循环，直接向量化操作即可：

import jax
import jax.numpy as jnp
from jax import jit

key = jax.random.PRNGKey(42)

@jit
def get_data():
  data = jax.random.normal(key, (5, 3))
  data = data.at[-2:].set(0.)
  return data

data = get_data()

@jit
def filter_nonzero(data):
  # 生成逐行的非零判断掩码
  row_mask = jnp.all(data != 0, axis=1)
  # 直接通过掩码取符合条件的行
  return data[row_mask]

accu = filter_nonzero(data)
print(accu.shape) # 输出(3, 3)，符合预期

方案二：通过carry传递累加状态（适合自定义复杂循环逻辑的场景）
如果你的筛选逻辑更复杂，必须用lax.while_loop实现，就把累加的数组作为循环的carry状态，在迭代间显式传递，最后作为返回值输出：

from jax import lax

@jit
def filter_with_loop(data):
  # 初始化循环carry：(当前遍历索引, 累加结果数组)，和你之前的初始逻辑对齐，初始值取第0行
  init_carry = (1, data[0:1])

  def cond_fun(carry):
    i, _ = carry
    return jnp.all(data[i])

  def body_fun(carry):
    i, acc = carry
    new_acc = jnp.vstack((acc, data[i]))
    return (i + 1, new_acc)

  _, accu = lax.while_loop(cond_fun, body_fun, init_carry)
  return accu

accu2 = filter_with_loop(data)
print(accu2.shape) # 输出(3, 3)，符合预期

注意：JAX中所有需要在控制流（lax.while_loop/lax.scan/lax.cond）之间传递的状态，都必须作为显式参数传递，不要依赖Python全局变量修改、列表append这类副作用操作，这类操作只会在追踪阶段执行一次，不会进入编译后的计算图。

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