DikobrAz
DikobrAz

Reputation: 3707

Stack vectors of different lengths in Tensorflow

How can I stack vectors of different length in tensorflow, e.g. from

[1, 3, 5]
[2, 3, 9, 1, 1]
[6, 2]

get zero-padded matrix

[1, 3, 5, 0, 0]
[2, 3, 9, 1, 1]
[6, 2, 0, 0, 0]

Vector count is known at definition time, but their lengths are not. Vectors are produced using tf.where(condition)

Upvotes: 2

Views: 2203

Answers (1)

kmario23
kmario23

Reputation: 61355

One way you can do this is like:

In [11]: v1 = [1, 3, 5]
In [12]: v2 = [2, 3, 9, 1, 1]
In [14]: v3 = [6, 2]

In [38]: max_len = max(len(v1), len(v2), len(v3))
In [39]: pad1 = [[0, max_len-len(v1)]]
In [40]: pad2 = [[0, max_len-len(v2)]]
In [41]: pad3 = [[0, max_len-len(v3)]]

# pads 0 to original vectors up to `max_len` length
In [42]: v1_padded = tf.pad(v1, pad1, mode='CONSTANT')
In [43]: v2_padded = tf.pad(v2, pad2, mode='CONSTANT')
In [44]: v3_padded = tf.pad(v3, pad3, mode='CONSTANT')


In [53]: res = tf.stack([v1_padded, v2_padded, v3_padded], axis=0)

In [56]: res.eval()
Out[56]: 
array([[1, 3, 5, 0, 0],
       [2, 3, 9, 1, 1],
       [6, 2, 0, 0, 0]], dtype=int32)

To make it work with N vectors efficiently, you should probably use a for loop to prepare the pad variables for all the vectors and the padded vectors subsequently. And, finally use tf.stack to stack these padded vectors along the 0th axis to get your desired result.


P.S.: You can get the length of the vectors dynamically once they are obtained from tf.where(condition).

Upvotes: 3

Related Questions