I’m looking for an elegant pytorch way to take a 2d matrix and move all non-zero values to the beginning of each row, while keeping their order. My current solution is super hacky, there must be something nicer?

For example:

input:

[[1, 5, 0, 2, 0, 0],

[0, 0, 9, 7, 0, 0]

output:

[[1, 5, 2, 0, 0, 0],

[9, 7, 0, 0, 0, 0]