使用布尔掩码从两个较小的张量创建 TensorFlow 张量

Question

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

说明我的意思 - 在 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)

简单说明：

使用布尔掩码从两个较小的张量创建 TensorFlow 张量

Create TensorFlow tensor from two smaller tensors using boolean mask

python

numpy

tensorflow