jax.lax.switch

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

根据 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]) – 基于 index 应用的函数序列 (A -> B)。所有分支必须返回相同的输出结构。

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

返回值:

基于 index 选择的分支的 branch(*operands) 的值 (B)。