How to get rows of matrix without using any loops in numpy or tensorflow?

Question

I have a matrix and I want to read in each row of the matrix and use tf.nn.top_k to find the top k values for each row.

How would I get each row of the matrix without using loops or list comprehension. I want to use Tensorflow or numpy and I was thinking of applying broadcasting.

If I put the indices in an array, for example if there were 10 rows in the matrix, I would have:

indices = [0,1,2,3......,9]

and then I could apply broadcasting?

Answer 1

You can do this a number of different ways. The following way uses numpy's argsort .

import numpy as np

a = np.random.randint(0, 1000, (4, 4))
a

returns

array([[712, 654, 871, 934],
       [639, 587, 216, 952],
       [853, 311, 598, 222],
       [466, 957, 306, 359]])

Then we can get the descending order of the arguments and use those to extract the corresponding elements:

o = np.argsort(a)[:, ::-1]
a[np.arange(4)[:, None], o[:, :2]]

returns

array([[934, 871],
       [952, 639],
       [853, 598],
       [957, 466]])

Answer 2

From the documentation of tf.nn.top_k

Finds values and indices of the k largest entries for the last dimension.

So, you can directly pass your tensor to this function and it will find the k largest elements from the last dimension. No need to use loops.

Answer 3

Say (like in one other answer here):

a = np.random.randint(0, 1000, (4,4))

then you can simply do:

np.sort(a)[:,-1:-3:-1]

where instead of the 3 index you put a number related to how many top values you want to get.

EDIT : In order to address your comment: First, I change to float datatype (to deal with the 1/3 value):

a = 1000*np.random.random((4,4))

then I get the indexes:

idx = a.argsort().argsort()

assign 1/3 to top two values:

a[idx>=2] = 1./3.

and assign zero to rest:

a[idx<2] = 0

et voila.

Answer 4

Here is a tensorflow way to do this.

#+BEGIN_SRC python :results output org drawer
import numpy as np
import tensorflow as tf

mat = np.random.randint(0, 1000, (4, 4))
print(mat)

# All tensorflow from here
tmat = tf.convert_to_tensor(mat)

k = 3

# ki are the indices within each row
kv, ki = tf.nn.top_k(tmat, k)
kmin = tf.reduce_min(kv, axis=1)[:, None]

newmat = tf.where(tmat < kmin, tf.zeros_like(tmat, dtype=tf.float64), tf.ones_like(tmat) / k)

with tf.Session() as sess:
    print(newmat.eval())
#+END_SRC

#+RESULTS:
:RESULTS:
[[ 97 805   0 104]
 [920 340 905  24]
 [439 243 705 161]
 [146 650 542 284]]
[[0.33333333 0.33333333 0.         0.33333333]
 [0.33333333 0.33333333 0.33333333 0.        ]
 [0.33333333 0.33333333 0.33333333 0.        ]
 [0.         0.33333333 0.33333333 0.33333333]]
:END: