Iteratively build tensor in Tensorflow

Question

Let's say I have a function that takes a Tensor in input (of a given dimensionality) and returns another Tensor in output. I would like to use that function on a batch of inputs and I would like it to return a batch of outputs. So both the input and the output would have one more dimension.

I could write a tf.while_loop to execute my function on all the inputs in the batch, but I am unsure on how to store the output of the single elements in the batch. I have an Idea on how to do this that should also clarify what I am trying to do, but I am not sure it would be optimal.

batch = tf.random.uniform([4,3,2]) #batch of size 4 of (3,2) shaped tensors
output = tf.zeros([0,5]) #let's say that the output should be a batch of 4 (4,5) shaped     tensors.
#I will concatenate the single outputs to this tensor and then reshape it
for i in tf.range(len(batch)):
 output = tf.concat((output,MyVeryNiceFunction(batch[i])),0) #MyVeryNiceFunction     returns a (4,5) shaped tensor
output = tf.reshape(output,(4,4,5)) #(batch_size,(shape of tensor))
return output 

This code for sure gives the output I want, but would it allow to parallelize each execution of the loop? Is there a better way to do this? Is there a proper data structure that would allow me to store the output for each loop execution, and then efficiently build the output Tensor from that?

Answer 1

In general, iterating over a dimension is very likely to be the wrong approach. In TF (and Matlab and Numpy), the goal is vectorization - describing your operations in a way that can touch all elements of the batch at the same time.

For example, let's say my dataset is composed of length 2 vectors, and I have a batch of 4 of them.

data = tf.convert_to_tensor([[1,2], [3,4], [5,6], [7,8]], tf.float32)
>>> data
<tf.Tensor: shape=(4, 2), dtype=float32, numpy=
array([[1., 2.],
       [3., 4.],
       [5., 6.],
       [7., 8.]], dtype=float32)>

If you wanted to add an element to each vector in a vectorized way, adding some kind of statistical analysis such as variance, you'd do this. Notice how you are constantly thinking about tensors shapes and dimensions and how to concat/append tensors. It's common to document tensor shapes constantly and even assert them. Welcome to TF programming.

vars = tf.math.reduce_variance(data, axis=1, keepdims=True)
tf.debugging.assert_equal(tf.shape(vars), [4, 1])
tf.concat(values=[data, vars], axis=1)


<tf.Tensor: shape=(4, 3), dtype=float32, numpy=
array([[1.  , 2.  , 0.25],
       [3.  , 4.  , 0.25],
       [5.  , 6.  , 0.25],
       [7.  , 8.  , 0.25]], dtype=float32)>