Smart indexing question

1

Hi, guys,

Let us have two tensors A and B. I want to obtain the tensor C of indices, containing the indices in A for the elements in B. The shape of C is the same as B.

For instance:

A = t.Tensor([1,2])
B = t.Tensor([[1,1,2],[2,2,1]])
C = (A == B)

I expect the result to be:

C = t.Tensor([[0,0,1],[1,1,0]])
2

Hey Zhenya,

I think something like this will do:

A = t.Tensor([1,2])
B = t.Tensor([[1,1,2],[2,2,1]])

C = torch.max(A.view(-1, 1, 1) == B.unsqueeze(0), dim=0).indices

Tell if that helped :slight_smile:

3

It helped, thank you!

1 Like