mindspore.ops.where
- mindspore.ops.where(condition, input, other)[源代码]
返回一个Tensor,Tensor的元素从 input 或 other 中根据 condition 选择。
\[\begin{split}output_i = \begin{cases} input_i,\quad &if\ condition_i \\ other_i,\quad &otherwise \end{cases}\end{split}\]- 参数:
condition (Tensor[bool]) - 如果是
True
,选取 input 中的元素,否则选取 other 中的元素。input (Union[Tensor, Scalar]) - 在 condition 为
True
的索引处选择的值。other (Union[Tensor, Scalar]) - 当 condition 为
False
的索引处选择的值。
- 返回:
Tensor,其中的元素从 input 和 other 中选取。
- 异常:
TypeError - 如果 condition 不是Tensor。
TypeError - 如果 input 和 other 都是常量。
ValueError - condition 、 input 和 other 不能互相广播。
- 支持平台:
Ascend
GPU
CPU
样例:
>>> import numpy as np >>> from mindspore import Tensor, ops >>> from mindspore import dtype as mstype >>> a = Tensor(np.arange(4).reshape((2, 2)), mstype.float32) >>> b = Tensor(np.ones((2, 2)), mstype.float32) >>> condition = a < 3 >>> output = ops.where(condition, a, b) >>> print(output) [[0. 1.] [2. 1.]]