TensorFlow: stacking tensors in while loop

Question

I'm trying to implement a loop that iterates over the rows of a tensor, retrieve the indices in each row, use them to gather vectors from another tensor and finally combine those vector in a new tensor. The problem is that each row may contain a different number of indices (e.g. [[-1,-1,1,4,-1], [3,-1,-1,-1,-1]] first row indices: [1, 4]; second row indices [3]). The problem rises when I use tf.while_loop or tf.scan. With the first one I don't understand how to stack all the gathered tensors together. The second one, instead, wants all the outputs to have the same shape (seems like i cannot tell that all the outputs have a general shape of [None, 10]).

Does anybody ever tried something similar?

I'm attaching the code for the while_loop:

i = tf.constant(0)
def body(i, merging):
    i += 1
    print('i', i)
    i_row = tf.gather(dense, [i])
    i_indices = tf.where(i_row > 0)[:, 1]
    i_vecs = tf.gather(embeddings_ph, i_indices)
    return i, i_vecs

tf.while_loop(lambda i, merging : tf.less(i, 2), body, 
              loop_vars=[i,merging], 
              shape_invariants=[i.get_shape(), 
                                tf.TensorShape((None, 3))],
              name='vecs_gathering')

What its missing here is to stack all the while_loop outputs (i_vec for each i) together in a new tensors.

Answer 1

Ok, got the inspiration from the rnn implementation. I modified my code as follows and now it works perfectly:

def body(i, outputs):
    i_row = tf.gather(dense, [i])
    i_indices = tf.where(i_row > 0)[:, 1]
    i_vecs = tf.gather(embeddings_ph, i_indices)
    outputs = outputs.write(i, i_vecs)
    i += 1
return i, outputs

outputs = tf.TensorArray(dtype=tf.float32, infer_shape=False, size=1, 
                     dynamic_size=True) 
_, outputs = tf.while_loop(lambda i, *_: tf.less(i, 3), body,[0,outputs])

outputs = outputs.concat()

I want also to stress the fact that you MUST reassign the value of the TensorArray when you perform a write (otherwise tf will complain a lot about the fact you are not using the array you declared)