mlx.core.where

where(condition: scalar | array, x: scalar | array, y: scalar | array, /, *, stream: None | Stream |Device = None) → array

根据 condition 选择 x 或 y 中的元素。

条件数组和输入数组必须形状相同或可以相互广播。

参数：

• condition (array) – 条件数组。

• x (array) – 条件为 True 时选择的输入。

• y (array) – 条件为 False 时选择的输入。

返回：

包含从 x 和 y 中选取的元素的输出。

返回类型：

array