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?
CodePudding user response:
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: