tensorflow: shift zeros to the end - tensorflow

Question

Given a tensor (with numbers >= 0) in tensorflow. I need to shift all zeros to the end of each line and remove columns that only include 0's.
E.g.
0 2 3 4
0 1 0 5
2 3 1 0
should be transformed to
2 3 4
1 5 0
2 3 1
Is there any nice way to do this in tensorflow? Btw, the order of the non-zero elements should be the same (no sorting).

Answer 1

Ragged tensor method
The best way
def rm_zeros(pred):
pred = tf.cast(pred, tf.float32)
# num_non_zero element in every row
num_non_zero = tf.count_nonzero(pred, -1) #[3 2 3]
# flat input and remove all zeros
flat_pred = tf.reshape(pred, [-1])
mask = tf.math.logical_not(tf.equal(flat_pred, tf.zeros_like(flat_pred)))
flat_pred_without_zero = tf.boolean_mask(flat_pred, mask) #[2. 3. 4. 1. 5. 2. 3. 1.]
# create a ragged tensor and change it to tensor, rows will be padded to max length
ragged_pred = tf.RaggedTensor.from_row_lengths(values=flat_pred_without_zero, row_lengths=num_non_zero)
paded_pred = ragged_pred.to_tensor(default_value=0.)
return paded_pred
a = tf.constant([[0, 2, 3, 4],[0, 1, 0, 5],[2, 3, 1, 0]])
print(rm_zeros(a))
output
tf.Tensor(
[[2. 3. 4.]
[1. 5. 0.]
[1. 2. 3.]], shape=(3, 3), dtype=float32)
Sorted method
If you don't mind the original data get sorted, the code below might be helpful. Although it's not the best solution.
The idea here is
1. change all zeros to infinity
2. sort the tensor
3. change all infinity back to zeros
4. slice the tensor to get minimal padding
def rm_zeros_sorted(input):
input = tf.cast(input, tf.float32)
# 1. change all zeros to infinity
zero_to_inf = tf.where(tf.equal(input, tf.zeros_like(input)), np.inf*tf.ones_like(input), input)
# 2. sort the tensor
input_sorted = tf.sort(zero_to_inf, axis=-1, direction='ASCENDING')
# 3. change all infinity back to zeros
inf_to_zero = tf.where(tf.math.is_inf(input_sorted), tf.zeros_like(input_sorted), input_sorted)
# 4. slice the tensor to get minimal padding
num_non_zero = tf.count_nonzero(inf_to_zero, -1)
max_non_zero = tf.reduce_max(num_non_zero)
remove_useless_zero = inf_to_zero[..., 0:max_non_zero]
return remove_useless_zero
a = tf.constant([[0, 2, 3, 4],[0, 1, 0, 5],[2, 3, 1, 0]])
print(rm_zeros_sorted(a))
output
tf.Tensor(
[[2. 3. 4.]
[1. 5. 0.]
[1. 2. 3.]], shape=(3, 3), dtype=float32)

Answer 2

The code below gets the trick done, although I'm sure that there are more elegant solutions possible and I'm curious to see those. The annoying part is that you have different amounts of zeros for each row.
a = tf.constant([[0, 2, 3, 4],[0, 1, 0, 5],[2, 3, 1, 0]])
boolean_mask = tf.logical_not(tf.equal(a, tf.zeros_like(a)))
# all the non-zero values in a flat tensor
non_zero_values = tf.gather_nd(a, tf.where(boolean_mask))
# number of non-zero values in each row
n_non_zero = tf.reduce_sum(tf.cast(boolean_mask, tf.int64), axis=-1)
# max number of non-zeros -> this will be the padding length
max_non_zero = tf.reduce_max(n_non_zero).numpy()
(Here it gets ugly)
# Split the tensor into flat tensors with the non-zero values of each row
rows = tf.split(non_zero_values, n_non_zero)
# Pad with zeros wherever necessary and recombine into a single tensor
tf.stack([tf.pad(r, paddings=[[0, max_non_zero - r.get_shape().as_list()[0]]]) for r in rows])
Produces the desired result:
<tf.Tensor: id=49, shape=(3, 3), dtype=int32, numpy=
array([[2, 3, 4],
[1, 5, 0],
[2, 3, 1]], dtype=int32)>

Answer 3

def shift_zeros(data, mask):
data_flat = tf.boolean_mask(data, mask)
nonzero_lens = tf.reduce_sum(tf.cast(mask, dtype=tf.int32), axis=-1)
nonzero_mask = tf.sequence_mask(nonzero_lens, maxlen=tf.shape(mask)[-1])
nonzero_data = tf.scatter_nd(tf.cast(tf.where(nonzero_mask), dtype=tf.int32), data_flat, shape=tf.shape(data))
return nonzero_data