jax.lax.select#

jax.lax.select(pred, on_true, on_false)[source]#

根据布尔谓词在两个分支之间进行选择。

包装了 XLA 的 Select 运算符。

通常，select() 会导致对两个分支都进行评估，尽管编译器可能会在可能的情况下省略计算。 对于一个通常只评估单个分支的类似函数，请参阅 cond()。

参数:

• pred (ArrayLike) – 布尔数组

• on_true (ArrayLike) – 数组，包含在 pred 为 True 时返回的条目。 必须与 pred 具有相同的形状，并与 on_false 具有相同的形状和 dtype。

• on_false (ArrayLike) – 数组，包含在 pred 为 False 时返回的条目。 必须与 pred 具有相同的形状，并与 on_true 具有相同的形状和 dtype。

返回值:

与 on_true 和 on_false 具有相同形状和 dtype 的数组。

返回类型:

result