TimZaman
TimZaman

Reputation: 2707

Sorting an Array in TensorFlow

Let's assume I have an array in TensorFlow:

[ 0.12300211,  0.51767069,  0.13886075,  0.55363625],
[ 0.47279349,  0.50432992,  0.48080254,  0.51576483],
[ 0.84347934,  0.44505221,  0.88839239,  0.48857492],
[ 0.93650454,  0.43652734,  0.96464157,  0.47236174], ..

I would like to sort this array by the third column. How do I do this? I am able to sort each column individually using tf.nn.top_k(), which gives me the sorted values and the respective indices. I could use the indices of this third column to reorder the others, but I cannot find a reordering Op.

Assuming I want to keep things in-graph (no Python shenanigans):

Upvotes: 15

Views: 9662

Answers (1)

keveman
keveman

Reputation: 8487

The following works :

a = tf.constant(...) # the array
reordered = tf.gather(a, tf.nn.top_k(a[:, 2], k=4).indices)

Upvotes: 12

Related Questions