How to implement 'np.where()' operation with libtorch

C++
1

Hi, i can get a matrix with numpy like this:

import numpy as np
label = np.random.randint(1,6,(2,3))
tmp = np.array(np.where(label == 3)).transpose((1, 0))

if label is array([[1, 2, 3], [3, 3, 4]]). then tmp will be array([[0, 2], [1, 0], [1, 1]], dtype=int64).

and i want to get the same result for libtorch. here is my code

auto ten = torch::randint(1,10,{2,3});
auto y = torch::full({2,3},3)
auto ind = torch::where(torch::operator==(ten,y), ten, y);

and ind’s data is 3 3 3 3 3 3 [Variable[CPUType]{2,3}]

where’s my fault?

2

Note that you are basically using (label==3).nonzero() in your numpy code, since you only provide the condition without x and y. From the docs:

Note
When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero() . Using nonzero directly should be preferred, as it behaves correctly for subclasses. The rest of this documentation covers only the case where all three arguments are provided.

That being said, let’s check the PyTorch code.
Your C++ code should be the equivalent to this Python code:

torch.where(
    torch.from_numpy(label) == 3,
    torch.from_numpy(label),
    torch.tensor([3])
)

If the condition is True, the values from the first tensor will be drawn, else from the second one.
Since you are comparing the tensor to 3, all True conditions will draw exactly the 3 from label. For all False return values in the condition, you will draw the corresponding value from a tensor full of 3s, which is why your output will contain all 3s.