Tensorflow: Reshape a tensor according to a boolean mask

Question

I have a 1D tensor of values: a = tf.constant([0.1, 0.2, 0.3, 0.4])

and a nD boolean mask: b = tf.constant([[1, 1, 0], [0, 1, 1]])

The total number of 1\'s in b ma