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

Question

假设我有一个大小为[batch，D]的二阶张量（矩阵），并且我想用Dd列元素交换元素的前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 columnsA = tf.Variable(tf.zeros((batch, D + d)))# set A[:,:D] to initial value# copy first d columns to the the last d columnsA = A[:,-d:].assign(A[:,:d])# use slicing to ignore the first d columnsA = A[:,d:]