Обратный порядок некоторых элементов в Tensorflow

Question

Скажите, что у меня есть тензор DATA формы (M, N, 2). У меня также есть другой тензор IND формы (N), состоящий из нулей и единиц.

Если IND(i)==1, то DATA(:,i,0) и DATA(:,i,1) должны поменяться местами. Если IND(i)==0 они не поменяются местами.

Как я могу это сделать? Я знаю, что это можно сделать через tf.gather_nd, но я понятия не имею, как.

Answer 1

Вот одно из возможных решений с tf.equal, tf.where, tf.scater_nd_update, tf.gather_nd и tf.reverse_v2:

data = tf.Variable([[[1, 2],
                     [2, 3],
                     [3, 4],
                     [4, 5],
                     [5, 6]]])  # shape=(1,5,2)

# reverse elements where ind is 1
ind = tf.constant([1, 0, 1, 0, 1])  # shape(5,)

cond = tf.where(tf.equal([ind], 1))
match_data = tf.gather_nd(data, cond)
rev_match_data = tf.reverse_v2(match_data, axis=[-1])
data = tf.scatter_nd_update(data, cond, rev_match_data)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(data))
    #[[[2 1]
    # [2 3]
    # [4 3]
    # [4 5]
    # [6 5]]]

Answer 2

Один из способов, который не использует tf.gather_ind, заключается в следующем.Идея состоит в том, чтобы создать DATA1, который представляет собой DATA со всеми возможными перестановками (то есть результатом перестановки, если IND был вектором 1 с), и использовать маски, чтобы выбрать правильные значения из данных или данных 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)