Given a TensorArray with a fixed size and entries with uniform shapes, I want to go to a Tensor containing the same values, simply by having the index dimension of the TensorArray as a regular axis.
TensorArrays have a method called "gather" which purportedly should do just that. And, in fact, the following example works:
array = tf.TensorArray(tf.int32, size=3)
array.write(0, 10)
array.write(1, 20)
array.write(2, 30)
gathered = array.gather([0, 1, 2])
"gathered" then yields the desired Tensor:
tf.Tensor([10 20 30], shape=(3,), dtype=int32)
Unfortunately, this stops working when wrapping it inside a tf.function, like so:
@tf.function
def func():
array = tf.TensorArray(tf.int32, size=3)
array.write(0, 10)
array.write(1, 20)
array.write(2, 30)
gathered = array.gather([0, 1, 2])
return gathered
tensor = func()
"tensor" then wrongly yields the following Tensor:
tf.Tensor([0 0 0], shape=(3,), dtype=int32)
Do you have an explanation for this, or can you suggest an alternative way to go from a TensorArray to a Tensor inside a tf.function?