How to rank a 2-d tensor along an axis in tensorflow?

Question

I am trying to get the ranks of a 2-d tensor in Tensorflow. I can do that in numpy using something like:

import numpy as np
from scipy.stats import rankdata

a = np.array([[0,3,2], [6,5,4]])
ranks = np.apply_along_axis(rankdata, 1, a)

And ranks is:

array([[ 1.,  3.,  2.],
       [ 3.,  2.,  1.]])

My question is how can I do this in tensorflow?

import tensorflow as tf
a = tf.constant([[0,3,2], [6,5,4]])
sess = tf.InteractiveSession()
ranks = magic_function(a)
ranks.eval()

Answer 1

tf.nn.top_k would work for you, although it has slightly different semantics. Please read the documentation to know how to use it for your case. But here is the snippet to solve your example :

sess = tf.InteractiveSession()
a = tf.constant(np.array([[0,3,2], [6,5,4]]))
# tf.nn.top_k sorts in ascending order, so negate to switch the sense
_, ranks = tf.nn.top_k(-a, 3)
# top_k outputs 0 based indices, so add 1 to get the same
# effect as rankdata
ranks = ranks + 1
sess.run(ranks)

# output :
# array([[1, 3, 2],
#        [3, 2, 1]], dtype=int32)