np. Where and torch Where usage

np.where

numpy.where() There are two uses ：

1. numpy.where(condition, x, y)
   Satisfy condition, Output x, Not satisfied with output y

condition = [[True, False], [True, True]]
x = [[1, 2], [3, 4]]
y = [[9, 8], [7, 6]]
np.where(condition, x, y)

>>>
[[1 8]
 [3 4]]

2. numy.where(condition)
   Only condiiton, Output meets condition The subscript of the element , The coordinates here are expressed in tuple Given in the form of , How many dimensions does the original array usually have , Output tuple It contains several arrays , Corresponding to the coordinates of each dimension of the qualified element

x = np.arange(12).reshape((2, 3, 2))
z = np.where(x > 7)

>>>
[[[ 0  1]
  [ 2  3]
  [ 4  5]]

 [[ 6  7]
  [ 8  9]
  [10 11]]]
(array([1, 1, 1, 1]), array([1, 1, 2, 2]), array([0, 1, 0, 1]))

torch.where

torch.where() The function is used to merge two according to certain rules tensor type

import torch

a = torch.randn(2, 3)
b = torch.ones_like(a)
com = torch.where(a > 0, a, b)

>>>
tensor([[-0.9728,  1.2321, -0.1471],
        [ 0.3736,  0.5832,  0.5332]])
tensor([[1.0000, 1.2321, 1.0000],
        [0.3736, 0.5832, 0.5332]])