Return last non-zero element of each zero-padded row in a 2D array

Say I have a 2D array in which each row is zero-padded to the maximum length:

[[1, 2, 0],
[3, 0, 0],
[4, 5, 6]]

And I also have the length of the rows before they were padded:

[2, 1, 3]

I want to return the last non-zero element of each row, i.e.,

[2, 3, 6]

Is there an elegant(no loops) to achieve this?

1 Like

Hi,

I think this can help:

#inputs
x = torch.tensor([[1, 2, 0],
                  [3, 0, 0],
                  [4, 5, 6]])
y = torch.tensor([2, 1, 3]) - 1

# creating one-hot mask
# borrowed from https://discuss.pytorch.org/t/convert-int-into-one-hot-format/507/4?u=nikronic
mask = torch.LongTensor(len(y), len(y))
mask.zero_()
mask.scatter_(1, y.view(-1, 1), 1)

# applying mask to zero non desired values
res = x*mask

# extract values that are non zero in each row
res[tuple(res.nonzero().T)]

But I am not sure there is any better solution or not, it just came to my mind. I hope it helps!

Bests

1 Like