Tensorflow Map_fn Tensorarray Has Inconsistent Shapes

I am playing around with the map_fn function, and noticed that it outputs a TensorArray, which should mean it is capable of outputting 'jagged' tensors (where the tensors on the in

Solution 1:

While the tf.map_fn() does use tf.TensorArray objects internally, and a tf.TensorArray can hold objects of different size, this program won't work as-is because tf.map_fn() converts its tf.TensorArray result back to a tf.Tensor by stacking the elements together, and it is this operation that fails.

You can however implement the tf.TensorArray-based using the lower-lever tf.while_loop() op instead:

lengths = tf.placeholder(tf.int32)
num_elems = tf.shape(lengths)[0]
init_array = tf.TensorArray(tf.float32, size=num_elems)

defloop_body(i, ta):
  return i + 1, ta.write(i, tf.random_normal((lengths[i],), 0, 1))

_, result_array = tf.while_loop(
    lambda i, ta: i < num_elems, loop_body, [0, init_array])

Solution 2:

Building upon mrry's answer, some more examples that can be run under TF2.x

import tensorflow as tf
# ================= example 1 ==================
num_elems = 5
init_array = tf.TensorArray(tf.float32, size=num_elems, infer_shape=False)
lengths = tf.range(0, 5)
defloop_body(i, ta):
  return i + 1, ta.write(i, tf.random.normal((lengths[i],), 0, 1))

_, result_array = tf.while_loop(
    lambda i, ta: i < num_elems, loop_body, [0, init_array])

for i inrange(num_elems):
    print(result_array.read(i))

# ================== example 2 ==================# TensorArray whose size is known at run time and shapes of elements# are not necessarily the same
ta = tf.TensorArray(tf.float32, size=0, dynamic_size=True, infer_shape=False)

# init ta with some mock data
ta = ta.write(0, 0.0)
ta = ta.write(1, 1.0)
ta = ta.write(2, tf.constant([2.0, 2.0]))

# loop bodydefloop_body(i, t):
    val = t.read(i)
    # do something
    t = t.write(i, tf.multiply(2.0, val))
    return i+1, t

# stop condition for while loop
index = tf.constant(0)
cond = lambda i, t: tf.less(i, t.size())

# results
i = tf.constant(0)
_, result_array = tf.while_loop(cond, loop_body, [i, ta])

for i inrange(result_array.size()):
    print(result_array.read(i))