CODEWITHSUNDEEP
pythontensorflow
severian
severian

Reputation: 498

Tensorflow: "staggered" sequence mask from indices?

Input like:

[1, 3, 2]

Desired output like (in a proper tensor):

[1 0 0 
 0 1 0
 0 1 0
 0 1 0
 0 0 1
 0 0 1]

I.e., very similar to tf.sequence_mask (which would give something like:

[1 1 1
 0 1 1
 0 1 0]

), but each subsequent element is "staggered" to start after the the prior sequence mask finishes.

Help greatly appreciated.

Upvotes: 2

Views: 84

Answers (1)

Vlad
Vlad

Reputation: 8585

It could be done by taking a square identity matrix of size equal to the number of elements in your input and then by applying tf.tile() inputs[i] number of times for each row i in the identity matrix:

import tensorflow as tf

inputs = tf.constant([1, 3, 2])

unit = tf.eye(num_rows=inputs.get_shape().as_list()[0])
unstacked = tf.unstack(unit)
tiled = [tf.tile(u[None, ...], multiples=[inputs[i], 1])
         for i, u in enumerate(unstacked)]
res = tf.concat(tiled, axis=0)

with tf.Session() as sess:
    print(sess.run(res))
# [[1. 0. 0.]
#  [0. 1. 0.]
#  [0. 1. 0.]
#  [0. 1. 0.]
#  [0. 0. 1.]
#  [0. 0. 1.]]

Upvotes: 3

Related Questions