jax.lax.scan

目录

jax.lax.scan#

jax.lax.scan(f, init, xs=None, length=None, reverse=False, unroll=1, _split_transpose=False)[源代码][源代码]#

在沿用状态的同时，扫描函数遍历数组的前导轴。

简而言之，Haskell-like type signature 是

scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])

对于任何数组类型说明符 t，[t] 表示在其前导轴上增加一个轴的类型，如果 t 是一个带有数组叶子的 pytree（容器）类型，那么 [t] 表示具有相同 pytree 结构且每个对应叶子在其前导轴上增加一个轴的类型。

当 xs 的类型（如上所述为 a）是数组类型或 None，并且 ys 的类型（如上所述为 b）是数组类型时，scan() 的语义大致由以下 Python 实现给出:

def scan(f, init, xs, length=None):
  if xs is None:
    xs = [None] * length
  carry = init
  ys = []
  for x in xs:
    carry, y = f(carry, x)
    ys.append(y)
  return carry, np.stack(ys)

与那个Python版本不同，xs 和 ys 可以是任意的 pytree 值，因此可以同时扫描多个数组并生成多个输出数组。None 实际上是这种情况的一个特例，因为它代表一个空的 pytree。

与那个Python版本不同，scan() 是一个JAX原语，并且被降低为一个单一的WhileOp。这使得它在减少JIT编译函数的编译时间方面很有用，因为在 jit() 函数中的原生Python循环结构会被展开，导致大型XLA计算。

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

备注

scan() 编译了 f，因此虽然它可以与 jit() 结合使用，但通常没有必要。

参数:

  • f (Callable[[Carry, X], tuple[Carry, Y]]) – 一个Python函数，类型为 c -> a -> (c, b)，这意味着 f 接受两个参数，其中第一个是循环进位的值，第二个是 xs 沿其主轴的切片，并且 f 返回一个对，其中第一个元素表示循环进位的新值，第二个元素表示输出的切片。

  • init (Carry) – 类型为 c 的初始循环进位值，可以是标量、数组或任何 pytree（嵌套的 Python 元组/列表/字典），表示初始循环进位值。此值必须与 f 返回的对中的第一个元素具有相同的结构。

  • xs (X | None) – 沿主要轴扫描的类型 [a] 的值，其中 [a] 可以是数组或任何具有一致主要轴大小的嵌套 Python 元组/列表/字典（pytree）。

  • length (int | None) – 可选的整数，指定循环迭代的次数，必须与 xs 中数组的前导轴的大小一致（但可用于不需要输入 xs 的扫描操作）。

  • reverse (bool) – 可选的布尔值，指定是否向前（默认）或反向运行扫描迭代，相当于反转 xs 和 ys 中数组的领先轴。

  • unroll (int | bool) – 可选的正整数或布尔值，指定在扫描原语的底层操作中，在一个循环迭代中展开多少次扫描迭代。如果提供一个整数，它决定了在一个循环的展开迭代中运行多少次展开的循环迭代。如果提供一个布尔值，它将决定循环是完全展开（即 unroll=True）还是完全不展开（即 unroll=False）。

  • _split_transpose (bool) – 实验性可选布尔值，指定是否进一步将转置拆分为扫描（计算激活梯度）和映射（计算与数组参数对应的梯度）。启用此功能可能会增加内存需求，因此这是一个可能会演变甚至被回滚的实验性功能。

返回:

类型为 (c, [b]) 的一对，其中第一个元素表示最终的循环进位值，第二个元素表示在输入的前导轴上扫描 f 的第二个输出时的堆叠输出。

返回类型:

tuple[Carry, Y]

目录