Tensorflow 中某些元素的倒序

Question

假设我有一个DATAshape的张量(M, N, 2)。我还有另一个IND形状为 (N) 的张量，由 0 和 1 组成。如果IND(i)==1那时DATA(:,i,0)和DATA(:,i,1)必须交换。如果IND(i)==0他们不换。我怎样才能做到这一点？我知道这可以通过 完成tf.gather_nd，但我不知道如何。

Answer 1

不使用 tf.gather_ind 的一种方法如下。这个想法是构建 DATA1，它是具有所有可能交换的 DATA（即交换的结果，如果 IND 是 1 的向量），并使用掩码从 Data 或 Data1 中选择正确的值，具体取决于是否需要交换或不。

DATA1 = tf.concat([tf.reshape(DATA[:,:,1], [M, N, 1]), tf.reshape(DATA[:,:,0], [M, N, 1])], axis = 2)

Mask1 = tf.cast(tf.reshape(IND, [1, N, 1]), tf.float64)

Mask0 = 1 - Mask1

Res = tf.multiply(Mask0, DATA) + tf.multiply(Mask1, DATA1)