Boolean mask Tensorflow, tf.boolean_mask - Maintain dimensions of original tensor

Question

I have a 1d tensor and a 1d array of Booleans of the same length.

I would like to use the Boolean array as a mask on the tensor, such that True keeps the original element-value in the tensor and False sets the original element-value in the tensor to zero.

E.g.

Tensor = [1,2,3,4,5]
Array = [True, False, False, False, True]

Apply Boolean mask to tensor:

Desired result = [1, 0, 0, 0, 5]
Result with tf.boolean_mask = [1, 5]

I have tried to use tf.boolean_mask(tensor, array), however, this reduces the dimensions of the resulting tensor to include only True elements, 2 dimensions in the above example.

How can I apply a Boolean mask to a tensor while maintaining the original dimensions of the tensor?

Answer 1

You can use tf.where:

tf.where(array, tensor, tf.zeros_like(tensor))

Answer 2

You can cast your boolean mask to a tf.int32 tensor and use tf.multiply:

mask = tf.constant([True, False, False, False, True])
A = tf.range(1,6)

with tf.Session() as sess:
   res = sess.run( \
          tf.multiply(A, tf.cast(mask, tf.int32)) \
         ) # [1, 0, 0, 0, 5]