jax.lax.switch#

jax.lax.switch(index, branches, *operands, operand=<object object>)[source]#

根据 index 应用 branches 中的恰好一个分支。

如果 index 超出范围，它将被钳制在有效范围内。

具有以下 Python 代码的语义

def switch(index, branches, *operands):
  index = clamp(0, index, len(branches) - 1)
  return branches[index](*operands)

在内部，这封装了 XLA 的 Conditional 运算符。然而，当通过 vmap() 转换以对一批谓词进行操作时，cond 会被转换为 select()。

参数:

  • index – 整数标量类型，指示要应用哪个分支函数。

  • branches (Sequence[Callable]) – 函数序列 (A -> B)，根据 index 应用。所有分支必须返回相同的输出结构。

  • operands – 应用的分支的输入操作数 (A)。

返回:

根据 index 选择的分支的 branch(*operands) 的值 (B)。