megengine.functional.where

where(mask, x=None, y=None)

Selects elements either from Tensor x or Tensor y, according to mask.

\[\textrm{out}_i = x_i \textrm{ if } \textrm{mask}_i \textrm{ is True else } y_i\]

Parameters

• mask (Tensor) – a mask used for choosing x or y.

• x (Optional[Tensor]) – first choice.

• y (Optional[Tensor]) – second choice.

Return type

Tensor

Returns

output tensor.

Examples

>>> import numpy as np
>>> mask = Tensor(np.array([[True, False], [False, True]], dtype=np.bool))
>>> x = Tensor(np.array([[1, np.inf], [np.nan, 4]],
...     dtype=np.float32))
>>> y = Tensor(np.array([[5, 6], [7, 8]], dtype=np.float32))
>>> out = F.where(mask, x, y)
>>> out.numpy()
array([[1., 6.],
       [7., 4.]], dtype=float32)

previous

megengine.functional.cond_take

next

megengine.functional.non_zero