Bubble up non null rows of a 3-D Tensor PyTorch
I am having some problem trying to convert the following function into a function only manipulating tensors.
def valid_sequence_output(sequence_output, valid_mask):
    bs, max_len, feat_dim = sequence_output.shape
    valid_output = torch.zeros(bs, max_len, feat_dim, dtype=torch.float32)
    for i in range(bs):
        jj = -1
        for j in range(max_len):
            if valid_mask[i][j].item() == 1:
                jj += 1
                valid_output[i][jj] = sequence_output[i][j]
    return valid_output

where the input tensors can be created as follow:
size = ((2,5,2))
sequence_output = torch.randint(0, 250, size=size)
valid_mask = torch.randint(0, 2, size=size[:2])

I basically aim at “bubbling up” the non null rows of sequence_output. As an example having sequence_output equal to:
tensor([[[  0,   0],
         [ 15,  47],
         [124, 230],
         [ 0,  0],
         [  65,   31]],
        [[  0,   0],
         [  0,   0],
         [  0,   0],
         [139, 228],
         [224, 205]]])

I am trying to obtain the following tensor:
tensor([[[ 15,  47],
         [124, 230],
         [ 65,  31],
         [  0,   0],
         [  0,   0]],
        [[139, 228],
         [224, 205],
         [  0,   0],
         [  0,   0],
         [  0,   0]]])

If somebody has a suggestion on how to do this, I would really appreciate :D

I managed to come up with a quite nasty solution (which he's probably suboptimal) by constructing a matrix A of zeros and ones that will swap rows of X when performing a matrix multiplication...
def vso(seq_out, valid_mask):
    X = torch.where(valid_mask.unsqueeze(-1) == 1, seq_out, torch.zeros_like(seq_out))
    bs, max_len, _ = X.shape
    tu = torch.unique(torch.nonzero(X)[:, :2], dim=0)
    batch_axis = tu[:, 0]
    rows_axis = tu[:, 1]
    a = torch.arange(bs).repeat(batch_axis.shape).reshape(batch_axis.shape[0], -1).T
    T = torch.cumsum(batch_axis == a, dim=1) - 1
    cols_axis = T[batch_axis, torch.arange(batch_axis.shape[0])]
    A = torch.zeros((bs, max_len, max_len))
    A[(batch_axis, cols_axis, rows_axis)] = 1
    valid_output = torch.matmul(A, X)
    return valid_output

Still looking for better answers though !
