Create TensorFlow tensor from two smaller tensors using boolean mask

Question

I'm using TensorFlow and would like to create a 1D tensor t1 from two smaller tensors t2 and t3, where len(t2) + len(t3) == len(t1) and a boolean mask which indicates how t2 and t3 should be combined. The boolean mask indicates how to "splice" the two smaller tensors together.

To show what I mean - in numpy, this is fairly easy:

mask = np.array([True, True, False, False, True])

a2 = [1., 2., 3.]
a3 = [4., 5.]

a1 = np.zeros(5)
a1[mask] = a2  # use mask to splice smaller arrays together
a1[~mask] = a3

a1  # array([1., 2., 4., 5., 3.])

I've had a look around and can't seem to find any equivalent code for TensorFlow. tf.where seems to require all arguments to be of the same size, which isn't possible in my use case. Is there a simple and moderately efficient way to do this?

Answer 1

Maybe try using tf.tensor_scatter_nd_update:

import tensorflow as tf

mask = tf.constant([True, True, False, False, True])

a2 = tf.constant([1., 2., 3.])
a3 = tf.constant([4., 5.])
a1 = tf.tensor_scatter_nd_update(tf.zeros((5,)), tf.concat([tf.where(mask),tf.where(~mask)], axis=0), tf.concat([a2, a3], axis=0))
print(a1)
# tf.Tensor([1. 2. 4. 5. 3.], shape=(5,), dtype=float32)

Simple explanation: