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 的条件运算符。 但是，当使用 vmap() 进行转换以对一批谓词进行操作时，cond 将转换为 select()。

参数:

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

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

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

返回值:

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