Tensor indexing with conditions

Hi everyone,
suppose I have a tensor: x: [0, 1, 4, 5, 6]
and I want to index elements in any of this list: [4, 5, 6],
similar to: np.where(x in [4, 5, 6], True, False) (this also doesn’t work :joy:),
Expected outcome: [False, False, True, True, True]
Are there any solutions? Thanks.

Could you explain the use case a bit more or give an example output for the provided tensors?
Would you like to select a random value from [4, 5, 6]?

Hi, I think I didn’t explain the example enough.
So the input x can be: [1, 2, 3, 4, 5, 6]
I want to locate the index of elements i in x where i also belongs to y:[4, 5, 6]
The expected outcome is: [False, False, False, True, True, True].

I managed to achieve this by numpy:
index = np.isin(x.numpy(), [4, 5, 6])

I am wondering if pytorch has similar functions. Thank you.

PyTorch doesn’t have the same isin implementation, but you might take a look at the numpy implementation and maybe use their method?

Alternatively, you could of course use a list comprehension as:

x = torch.arange(7)
y = torch.tensor([4, 5, 6])

ret = torch.stack([y_ == x for y_ in y]).sum(0).bool()

Yes, .sum(0) is a smart trick. Thank you