在带有 Tensorflow 的 Keras 中，如何重新索引 nd 张量的轴？

Question

像这样，除了没有错误：

input = tf.convert_to_tensor(np.random.rand(500,100,5))
new_order = [0,4,1,3,2]
output = input[:,:,new_order]

我发现最接近的是tf.gather ，但我无法让它工作。

Answer 1

我发现很容易转置张量，以便索引维度是第一个维度，收集元素然后将其放回原始形状。

output = tf.transpose(
    tf.gather(
        tf.transpose(input, [2,0,1]), 
        new_order
    ), [1,2,0]
)