Change the order of elements in first axis tensorflow

Question

I have an n dimensional tensor . I want to change the order of elements based on first axis and a given order. For example (1, 0, 2, 4, 3, 5) should give me this result for this matrix:

[1, 0, 0, 0, 0]        [0, 1, 0, 0, 0]     
[0, 1, 0, 0, 0]        [1, 0, 0, 0, 0]
[0, 0, 1, 0, 0]   ==>  [0, 0, 1, 0, 0]
[0, 0, 0, 1, 0]        [0, 0, 0, 0, 1]
[0, 0, 0, 0, 1]        [0, 0, 0, 1, 0]
[0, 0, 0, 0, 1]        [0, 0, 0, 0, 1]

The order is important for me because I have some tensors and I want all of them to be reordered with the same order. How can I achieve this?

Answer 1

You should use tf.gather:

with tf.Session() as sess:
    data = [[1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0],
            [0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1],
            [0, 0, 0, 0, 1]]
    idx = [1, 0, 2, 4, 3, 5]
    result = tf.gather(data, idx)
    print(sess.run(result))

Output:

[[0 1 0 0 0]
 [1 0 0 0 0]
 [0 0 1 0 0]
 [0 0 0 0 1]
 [0 0 0 1 0]
 [0 0 0 0 1]]