用张量反求张量中的两个元素块

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:]