torch.where

torch.where(condition, x, y) → Tensor

Возвращает тензор элементов, выбранных из x или y, в зависимости от condition.

Операция определяется следующим образом:

$$\text{out}_i = \begin{cases} \text{x}_i & \text{if } \text{condition}_i \\ \text{y}_i & \text{otherwise} \\ \end{cases}$$

Примечание

Тензоры condition, x, y должны быть совместимы по правилам векторизации.

Параметры:

• condition (BoolTensor) – При True (ненулевое значение), возвращается x, иначе возвращается y
• x (Tensor или Скаляр) – значение (если x является скаляром) или значения, выбранные в индексах, где condition принимает значение True
• y (Tensor или Скаляр) – значение (если y является скаляром) или значения, выбранные в индексах, где condition принимает значение False

Возвращает:

Тензор с формой, равной результату векторизации condition, x, y

Тип возвращаемого значения:

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, 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) → кортеж LongTensor

torch.where(condition) идентично torch.nonzero(condition, as_tuple=True).

Примечание

См. также torch.nonzero().