How TensorArray and while_loop work together in tensorflow?

python, tensorflow

Solution

There are several things in your code to point out. First, you don't need to unstack the matrix into the `TensorArray` to use it inside the loop, you can safely reference the matrix `Tensor` inside the body and index it using `matrix[i]` notation. Another issue is the different data type between your matrix (`tf.int32`) and the `TensorArray` (`tf.float32`), based on your code you're multiplying the matrix ints by 2 and writing the result into the array so it should be int32 as well. Finally, when you wish to read the final result of the loop, the correct operation is `TensorArray.stack()` which is what you need to run in your `session.run` call.

Here's a working example:

import numpy as np
import tensorflow as tf    

# 1000 sequence in the length of 100
matrix = tf.placeholder(tf.int32, shape=(100, 1000), name="input_matrix")
matrix_rows = tf.shape(matrix)[0]
ta = tf.TensorArray(dtype=tf.int32, size=matrix_rows)

init_state = (0, ta)
condition = lambda i, _: i < matrix_rows
body = lambda i, ta: (i + 1, ta.write(i, matrix[i] * 2))
n, ta_final = tf.while_loop(condition, body, init_state)
# get the final result
ta_final_result = ta_final.stack()

# run the graph
with tf.Session() as sess:
    # print the output of ta_final_result
    print sess.run(ta_final_result, feed_dict={matrix: np.ones(shape=(100,1000), dtype=np.int32)}) 

Problem

Original source