python - 使用 bool 掩码从两个较小的张量创建 TensorFlow 张量

Question

我正在使用 TensorFlow 并想创建一维张量 t1来自两个较小的张量 t2和 t3 , 其中len(t2) + len(t3) == len(t1)和一个 bool 掩码，指示如何 t2和 t3应该合并。 bool 掩码指示如何将两个较小的张量“拼接”在一起。

为了说明我的意思——在 numpy 中，这相当简单:

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.])

我环顾四周，似乎找不到任何适用于 TensorFlow 的等效代码。 tf.where似乎要求所有参数都具有相同的大小，这在我的用例中是不可能的。有没有一种简单且效率适中的方法来做到这一点？

Answer 1

也许尝试使用 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)

简单解释: