jax.lax.while_loop

jax.lax.while_loop(cond_fun, body_fun, init_val)

当 cond_fun 为 True 时，重复调用 body_fun。

Haskell 风格的类型签名简述如下

while_loop :: (a -> Bool) -> (a -> a) -> a -> a

while_loop 的语义由此 Python 实现给出

def while_loop(cond_fun, body_fun, init_val):
  val = init_val
  while cond_fun(val):
    val = body_fun(val)
  return val

与 Python 版本不同，while_loop 是一个 JAX 原语，被降低为单个 WhileOp。这使其对于减少 jit 编译函数的编译时间非常有用，因为 @jit 函数中的原生 Python 循环结构会被展开，从而导致大型 XLA 计算。

与 Python 类似物不同，循环携带值 val 必须在所有迭代中保持固定的形状和 dtype（而不仅仅是在 NumPy 秩/形状广播和 dtype 提升规则中保持一致）。换句话说，上面类型签名中的类型 a 表示具有固定形状和 dtype 的数组（或具有固定结构和叶子处具有固定形状和 dtype 的数组的嵌套元组/列表/字典容器数据结构）。

与使用原生 Python 循环结构的不同之处还在于，while_loop 不是反向模式可微的，因为 XLA 计算需要在内存需求上具有静态边界。

注意

while_loop() 编译 cond_fun 和 body_fun，因此虽然它可以与 jit() 结合使用，但通常是不必要的。

参数:

• cond_fun (Callable[[T], BooleanNumeric]) – 类型为 a -> Bool 的函数。

• body_fun (Callable[[T], T]) – 类型为 a -> a 的函数。

• init_val (T) – 类型为 a 的值，它可以是标量、数组或任何 pytree（嵌套的 Python 元组/列表/字典），表示初始循环携带值。

返回:

body_fun 最后一次迭代的输出，类型为 a。

返回类型:

T