Pytorch efficient indexing a tensor over duplications by an entry value tensor

user_9125078 0
Viewed 130

Here is my user case:

A: tensor([[0,0,1,1,2,3],[5,6,7,8,9,3]]), # an input tensor 
B: tensor([2,0]), # an indexing tensor 
is there an efficient way to get the resulting tensor as:
O: tensor([[2,0,0], [9,5,6]]).

That is, B values match with A[0] entries, then indexing along columns. Note that B may be pre-sorted in ascending order (if make the operation simpler) such that the resulting tensor will be

O: tensor([[0,0,2],[5,6,9]]).

Thank you

pytorch indexing python
2 Answers
  1. Iterate over A[0] and get index of values that matches with values within B
  2. Iterate over A and get values located at computed indexes.
A = np.array([[0,0,1,1,2,3],[5,6,7,8,9,3]])
B = np.array([2,0])

indexes = [i for i,j in enumerate(A[0]) if j in B]
output = [list(i[indexes]) for i in A]
print(output)
> [[0, 0, 2], [5, 6, 9]]
user_135770630

Here is a solution I came up with:

class PairedValGather():
    def __call__(self, source, value):
        mask = self.make_mask(source, value)
        selected = torch.masked_select(source, mask)
        selected = selected.reshape(-1, selected.size()[0]//2)
        return selected

    def make_mask(self, source, ivalue):
        mask = torch.zeros(source.shape, dtype=torch.int64)
        for i, entry in enumerate(source[0]):
            item = entry.data.item()
            if item in ivalue[0] and source[1,i] in ivalue[1]:
                mask[:,i] = 1
        return mask>0

Usage:

source=torch.from_numpy(np.array([[1,1,2,3,4],[8,9,5,4,3]]))
ivalue=[[1,3],[9,4]]
selected = PairedValGather()(source, ivalue)
selected

Result:

tensor([[1, 3], [9, 4]], dtype=torch.int32)

I haven't done a performance measure but expect it good for a large source tensor. Any alternative better solution is still appreciated.

user_91250780
Related