In Keras with Tensorflow, how can I reindex an axis of nd Tensor?

Question

Like this, except without the errors:

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

Closest I found is tf.gather , but I haven't been able to make it work.

Answer 1

I find it easy to transpose the tensor so that the indexing dimension is the first dimension, gather the element and then put it back to the original shape.

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