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

• operands – 操作数 (A)，作为输入传递给所应用的任何分支。

返回:

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