CODEWITHSUNDEEP
pythonnumpytensorflowranking
Zichen Wang
Zichen Wang

Reputation: 1374

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

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()

Upvotes: 2

Views: 776

Answers (1)

keveman
keveman

Reputation: 8487

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)

Upvotes: 4

Related Questions