mindspore.mint.where

mindspore.mint.where(condition, input, other)[源代码]

返回一个Tensor,Tensor的元素从 inputother 中根据 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]) - 在 conditionTrue 的索引处选择的值。

  • other (Union[Tensor, Scalar]) - 当 conditionFalse 的索引处选择的值。

返回:

Tensor,其中的元素从 inputother 中选取。

异常:
  • TypeError - 如果 condition 不是Tensor。

  • TypeError - 如果 inputother 都是常量。

  • ValueError - conditioninputother 不能互相广播。

支持平台:

Ascend

样例:

>>> import numpy as np
>>> from mindspore import Tensor, mint
>>> 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 = mint.where(condition, a, b)
>>> print(output)
[[0. 1.]
 [2. 1.]]