jax.lax.fori_loop

内容

jax.lax.fori_loop#

jax.lax.fori_loop(lower, upper, body_fun, init_val, *, unroll=None)[source]#

从 lower 循环到 upper，通过降维到 jax.lax.while_loop()。

简而言之，Haskell 式类型签名 为

fori_loop :: Int -> Int -> ((Int, a) -> a) -> a -> a

The semantics of fori_loop are given by this Python implementation

def fori_loop(lower, upper, body_fun, init_val):
  val = init_val
  for i in range(lower, upper):
    val = body_fun(i, val)
  return val

如 Python 版本所示，设置 upper <= lower 将不会产生任何迭代。不支持负数或自定义增量。

与 Python 版本不同，fori_loop 是通过调用 jax.lax.while_loop() 或 jax.lax.scan() 来实现的。如果循环次数是静态的（意味着在追踪时已知，可能是因为 lower 和 upper 是 Python 整数字面量），则 fori_loop 是基于 scan() 实现的，并且支持反向模式自动微分；否则，将使用 while_loop，并且不支持反向模式自动微分。有关更多信息，请参阅这些函数的文档字符串。

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

注意

fori_loop() 会编译 body_fun，因此虽然它可以与 jit() 组合使用，但通常没有必要。

参数：

  • lower – 表示循环索引下界的整数（包含）

  • upper – 表示循环索引上界的整数（不包含）

  • body_fun – 类型为 (int, a) -> a 的函数。

  • init_val – 类型为 a 的初始循环携带值。

  • unroll (int | bool | None) – 一个可选的整数或布尔值，用于确定循环展开的程度。如果提供整数，则确定在循环的单个卷积迭代中运行多少个展开的循环迭代。如果提供布尔值，则它将确定循环是否完全展开（即 unroll=True）或完全不展开（即 unroll=False）。此参数仅在循环边界静态已知时适用。

返回值：

来自最终迭代的循环值，类型为 a。

内容