torch.where

torch.where(condition, input, other, *, out=None) → Tensor

根據 condition，傳回從 input 或 other 中選取的元素張量。

此運算定義為：

$$\text{out}_i = \begin{cases} \text{input}_i & \text{if } \text{condition}_i \\ \text{other}_i & \text{otherwise} \\ \end{cases}$$

注意

張量 condition、input、other 必須是可廣播的 (broadcastable)。

參數

• condition (BoolTensor) – 當為 True (非零) 時，產生 input，否則產生 other

• input (Tensor 或 Scalar) – 值 (如果 input 是一個純量) 或在 condition 為 True 的索引處選取的值

• other (Tensor 或 Scalar) – 值 (如果 other 是一個純量) 或在 condition 為 False 的索引處選取的值

關鍵字參數

out (Tensor, optional) – 輸出張量。

返回值

形狀等於 condition、input、other 的廣播形狀 (broadcasted shape) 的張量

返回值類型

Tensor

範例

>>> x = torch.randn(3, 2)
>>> y = torch.ones(3, 2)
>>> x
tensor([[-0.4620,  0.3139],
        [ 0.3898, -0.7197],
        [ 0.0478, -0.1657]])
>>> torch.where(x > 0, 1.0, 0.0)
tensor([[0., 1.],
        [1., 0.],
        [1., 0.]])
>>> torch.where(x > 0, x, y)
tensor([[ 1.0000,  0.3139],
        [ 0.3898,  1.0000],
        [ 0.0478,  1.0000]])
>>> x = torch.randn(2, 2, dtype=torch.double)
>>> x
tensor([[ 1.0779,  0.0383],
        [-0.8785, -1.1089]], dtype=torch.float64)
>>> torch.where(x > 0, x, 0.)
tensor([[1.0779, 0.0383],
        [0.0000, 0.0000]], dtype=torch.float64)

torch.where(condition) → tuple of LongTensor

torch.where(condition) 與 torch.nonzero(condition, as_tuple=True) 相同。

注意

另請參閱 torch.nonzero()。