[FIXED] how to get sub 2d tensor with specific value condition in pytorch?

Issue

I want to copy a 2-d torch tensor to a destination tensor containing only values until the first occurrence of 202 value and zero for the rest of the items like this:

source_t=tensor[[101,2001,2034,1045,202,3454,3453,1234,202]
         ,[101,1999,2808,202,17658,3454,202,0,0]
         ,[101,2012,3832,4027,3454,202,3454,9987,202]]

destination_t=tensor[[101,2001,2034,1045,202,0,0,0,0]
                    ,[101,1999,2808,202,0,0,0,0,0]
                    ,[101,2012,3832,4027,3454,202,0,0,0]]

how can I do it?

Solution

I make working and pretty efficient solution.

I have made a little bit more complex source tensor with additional rows with 202 in different places:

import copy
import torch

source_t = torch.tensor([[101, 2001, 2034, 1045, 202, 3454, 3453, 1234, 202],
                         [101, 1999, 2808, 202, 17658, 3454, 202, 0, 0],
                         [101, 2012, 3832, 4027, 3454, 2020, 3454, 9987, 2020],
                         [101, 2012, 3832, 4027, 3454, 202, 3454, 9987, 202],
                         [101, 2012, 3832, 4027, 3454, 2020, 3454, 9987, 202]
                         ])

At the start, we should find occurrences of the first 202. We can find all occurrences and then choose the first one:

index_202 = (source_t == 202).nonzero(as_tuple=False).numpy()
rows_for_replace = list()
columns_to_replace = list()

elements = source_t.shape[1]
current_ind = 0

while current_ind < len(index_202)-1:
    current = index_202[current_ind]
    element_ind = current[1] + 1
    rows_for_replace.extend([current[0]]*(elements-element_ind))
    while element_ind < elements:
        columns_to_replace.append(element_ind)
        element_ind += 1
    if current[0] == index_202[current_ind+1][0]:
        current_ind += 1
    current_ind += 1

After this operation, we have all our indexes which we should replace with zeros. 4 elements in the first row, 5 in the second, 3 in the fourth row and nothing in the third and fifth. rows_for_replace, columns_to_replace

([0, 0, 0, 0, 1, 1, 1, 1, 1, 3, 3, 3], [5, 5, 5, 5, 4, 4, 4, 4, 4, 6, 6, 6])

Then we just copy our source tensor and set new values in place:

destination_t = copy.deepcopy(source_t)
destination_t[rows_for_replace, columns_to_replace] = 0

Summary: source_t

tensor([[  101,  2001,  2034,  1045,   202,  3454,  3453,  1234,   202],
        [  101,  1999,  2808,   202, 17658,  3454,   202,     0,     0],
        [  101,  2012,  3832,  4027,  3454,  2020,  3454,  9987,  2020],
        [  101,  2012,  3832,  4027,  3454,   202,  3454,  9987,   202],
        [  101,  2012,  3832,  4027,  3454,  2020,  3454,  9987,   202]])

destination_t

tensor([[ 101, 2001, 2034, 1045,  202,    0,    0,    0,    0],
        [ 101, 1999, 2808,  202,    0,    0,    0,    0,    0],
        [ 101, 2012, 3832, 4027, 3454, 2020, 3454, 9987, 2020],
        [ 101, 2012, 3832, 4027, 3454,  202,    0,    0,    0],
        [ 101, 2012, 3832, 4027, 3454, 2020, 3454, 9987,  202]])

Answered By - Valentina Biryukova