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}\]

Warning

Starting after version 2.9.0, a single-argument overload mindspore.ops.where(condition) will be supported, which returns the coordinates of non-zero elements, equivalent to mindspore.ops.argwhere(condition). This change is backward-compatible; the three-argument call remains unchanged.

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