mindspore.ops.where

mindspore.ops.where(condition, input, other)[source]

Return a tensor in which the elements are selected from input or other based on the condition.

Support broadcasting.

\[\begin{split}output_i = \begin{cases} input_i,\quad &if\ condition_i \\ other_i,\quad &otherwise \end{cases}\end{split}\]

Parameters

condition (Tensor[bool]) – If True, yield input, otherwise yield other.
input (Union[Tensor, Scalar]) – When condition is True, values to select from.
other (Union[Tensor, Scalar]) – When condition is False, values to select from.

Returns

Tensor

Supported Platforms:: Ascend GPU CPU

Examples

>>> import mindspore
>>> input = mindspore.tensor([[0, 1],
...                           [2, 3]])
>>> other = mindspore.tensor([[1, 1],
...                           [1, 1]])
>>> condition = input < 3
>>> mindspore.ops.where(condition, input, other)
Tensor(shape=[2, 2], dtype=Int64, value=
[[0, 1],
 [2, 1]])