使用TensorFlow反转张量内的两个元素块

Question

假设我有一个大小为[batch，D]的秩2张量（矩阵），并且我想用D-d列元素交换元素的前d列的位置。我如何在TensorFlow中有效地做到这一点？  示例：一个大小为[4,4]的矩阵，我想将元素的前d = 2列的位置与左边的列块交换。

输入：

1 2 3 4
1 2 3 4
1 2 3 4
1 2 3 4

输出：

3 4 1 2
3 4 1 2
3 4 1 2
3 4 1 2

谢谢。

Answer 1

假设您的矩阵为A，那么tf.concat((A[:,d:], A[:,:d]), axis=1)应该给您您想要的东西。我不知道有任何直接的方法可以使它更高效，但是除非它成为您模型中的瓶颈，否则我不会担心。

更新：如果d << D并且您只需要执行一次操作，则由于减少了复制开销（但也使用了更多的内存），因此以下操作可能会更快

# allocate D + d columns
A = tf.Variable(tf.zeros((batch, D + d)))

# set A[:,:D] to initial value

# copy first d columns to the the last d columns
A = A[:,-d:].assign(A[:,:d])
# use slicing to ignore the first d columns
A = A[:,d:]