CODEWITHSUNDEEP
tensorflow
SRobertJames
SRobertJames

Reputation: 9253

TensorFlow Selecting entries (from one of two tensors) based on a boolean mask

I have three tensors, a, b, and mask, all of the same shape. I'd like to produce a new tensor c, such that each entry of c is taken from the corresponding entry of a iff the corresponding entry of mask is True; else, it is taken from the corresponding entry of b.

Example:

a = [0, 1, 2]
b = [10, 20, 30]
mask = [True, False, True]
c = [0, 20, 2]

How can I do this?

Upvotes: 1

Views: 1402

Answers (2)

Drag0
Drag0

Reputation: 8928

You can do it like this:

1) convert mask to ints (0 for false, 1 for true)
2) do element wise multiplication of int_mask with tensor 'a' 
    (elements that should not be included are going to be 0)
3) do logical_not on mask
4) convert logical_not_int_mask to ints 
   (again 0 for false, 1 for true values)
5) now just do element wise multiplication of logical_not_int_mask with tensor 'b' 
   (elements that should not be included are going to be 0)
6) Add tensors 'a' and 'b' together and there you have it.

In code it should look something like this:

# tensor 'a' is [0, 1, 2]
# tensor 'b' is [10, 20, 30]
# tensor 'mask' is [True, False, True]

int_mask = tf.cast(mask, tf.int32)
# Leave only important elements in 'a'
a = tf.mul(a, int_mask) 
mask = tf.logical_not(mask)
int_mask = tf.cast(mask, tf.int32)
b = tf.mul(b, int_mask)
result = tf.add(a, b)

Or simply use tf.select() function just like someone already mentioned.

Upvotes: 1

Yahia Zakaria
Yahia Zakaria

Reputation: 1206

Why not use tf.select(condition, t, e, name=None)

for your example:

c = tf.select(mask, a, b)

for more details about tf.select, visit Tensorflow Control Flow Documentation

Upvotes: 5

Related Questions