Get mask from tensor of indices?

Question

If I have a tensor of shape "batch_size * length * hidden_size", and I have another index span tensor of shape "batch_size * 2", where the index span tensor indicates the beginning and ending indices I want to select from the first tensor. Say if the index span tensor has values

[[1,3], [2, 4]]

then I want to get the following mask from the first tensor

[[0, 1, 1, 1, 0, 0, ...], [0, 0, 1, 1, 1, 0, 0, ...]]

Is there a way to do it with Tensorflow?

vijay m vijay m · Accepted Answer · 2018-07-12T16:35:43

We can obtain the above my using a combination of range (to obtain the indices) and tf.sparse_to_dense (to populate the ones):

batch_size = 2
length_hidden_size = 10

# Use range to populate the indices
r = tf.map_fn(lambda x: tf.range(x[0], x[1]+1),a)

# convert to full indices
mesh = tf.meshgrid(tf.range(tf.shape(r)[1]), tf.range(tf.shape(r)[0]))[1]
full_indices = tf.reshape(tf.stack([mesh, r], axis=2), [-1,2])

# Set 1s to the full indices
val = tf.ones(tf.shape(full_indices)[0])
dense = tf.sparse_to_dense(full_indices,tf.constant([batch_size,length_hidden_size]), val, validate_indices=False)

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

Get mask from tensor of indices?

1 Answers