mindspore.ops.where

View Source On Gitee
mindspore.ops.where(condition, x, y)[source]

Selects elements from x or y based on condition and returns a tensor.

\[\begin{split}output_i = \begin{cases} x_i,\quad &if\ condition_i \\ y_i,\quad &otherwise \end{cases}\end{split}\]
Parameters

  • condition (Tensor[bool]) – If True, yield x, otherwise yield y.

  • x (Union[Tensor, Scalar]) – When condition is True, values to select from.

  • y (Union[Tensor, Scalar]) – When condition is False, values to select from.

Returns

Tensor, elements are selected from x and y.

Raises

  • TypeError – If condition is not a Tensor.

  • TypeError – If both x and y are scalars.

  • ValueError – If condition, x and y can not broadcast to each other.

Supported Platforms:

Ascend GPU CPU

Examples

>>> 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.]]