Harry
Harry

Reputation: 103

Tensorflow: create boolean matrix given True indices

Suppose I have below two vectors, one for start index, the other for end index

start_index = [1, 1, 2]
end_index = [3, 4, 3]

and below I have the shape of the final boolean matrix

shape = [3, 6]

I want to generate below boolean matrix

bool_mat = [[False, True,  True, True, False, False]
            [False, True,  True, True, True,  False]
            [False, False, True, True, False, False]]

for each row True starts from index in start_index and end at index in end_index, and False elsewhere,

bool_mat[i, start_index[i]:end_index[i]+1] = True

How to do this in TensorFlow? thanks!

Upvotes: 1

Views: 599

Answers (1)

javidcf
javidcf

Reputation: 59701

You can do that like this:

import tensorflow as tf

start_index = tf.constant([1, 1, 2])
end_index = tf.constant([3, 4, 3])
shape = tf.constant([3, 6])
col = tf.range(shape[1])
result = (col >= start_index[:, tf.newaxis]) & (col <= end_index[:, tf.newaxis])
with tf.Session() as sess:
    print(sess.run(result))

Output:

[[False  True  True  True False False]
 [False  True  True  True  True False]
 [False False  True  True False False]]

Upvotes: 1

Related Questions