How can I move all zeroes to end of array

#1

I have a variable that looks like
Input: [[0,2,3,0,0,0,1,2],
[0,0,1,3,0,0,5,0]]

and I want to end up with the following variable
Output: [[2,3,1,2,0,0,0,0],
[1,3,5,0,0,0,0,0]]

I have tried the a[a.nonzero()] if a is the input, but it only works in 1xn dims
Is there a way to do this in PyTorch?

#2

You can do it this way:

import torch
x = torch.Tensor([[0,2,3,0,0,0,1,2], [0,0,1,3,0,0,5,0]])
y = torch.empty(0, x.size(1))
for r in x:
    nz = r.nonzero().squeeze()
    z = torch.zeros(r.numel() - nz.numel())
    z = torch.cat((r[nz], z)).unsqueeze(0)
    y = torch.cat((y, z))

print(y)

I’m not sure it’s the best way to do it though, but hopefully it could be good enough for your case.