How to get argmax for indices in a matrix/tensor?
Is there any way to perform a top-k operation on a matrix or tensor so that the relevant indices are returned?
For example:
&gt;&gt;&gt; import torch
&gt;&gt;&gt; matrix = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
&gt;&gt;&gt; print(matrix)
tensor([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
&gt;&gt;&gt; indices = max_index_function(matrix)
&gt;&gt;&gt; print(indices)
tensor([2, 2])

Something like this. If I use torch.argmax it's a little tricky to get what I want. I want to be able to effectively return the indices because I want to extend an argmax to a top-k function as well.
Right now I have something like this:
&gt;&gt;&gt; column_max_idx = torch.unique(torch.argmax(matrix, dim=0))
&gt;&gt;&gt; row_max_idx = torch.unique(torch.argmax(matrix, dim=1))
&gt;&gt;&gt; idx_pairs = torch.cartesian_prod(column_max_idx, row_max_idx)

It does what I want but I'm sure there's a better way. Thanks!
Edit
Another thing I've tried is using two for loops and sorting, but I also am wondering if there's a function I could use without loops.
&gt;&gt;&gt; idxs_and_values = []
&gt;&gt;&gt; for col_idx, col_element in enumerate(matrix):
...     for row_idx, row_element in enumerate(col_element):
...         idxs_and_values.append(((col_idx, row_idx), row_element.item())
&gt;&gt;&gt; sorted(idxs_and_values, key=lambda x: x[-1], reverse=True)
[((2, 2), 9), ((2, 1), 8), ((2, 0), 7), ((1, 2), 6), ((1, 1), 5), \
  ((1, 0), 4), ((0, 2), 3), ((0, 1), 2), ((0, 0), 1)]


how about topk?
import torch 
matrix = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
values, indices = matrix.flatten().topk(k=3)
print(values) 
# tensor([9, 8, 7])
print(indices)
# tensor([8, 7, 6])

Note that the indices are now pointing to the flattened vector. You can recover the original d-dimensional indexing with numpy's unravel_index:
indices = [np.unravel_index(i, matrix.shape) for i in indices]
print(indices)
# [(2, 2), (2, 1), (2, 0)]

Nicely wrapped in a function for any tensor order:
from torch import Tensor 
from numpy import unravel_index

def mytopk(x: Tensor, k: int) -&gt; list[tuple[int,...]]:
    values, indices = x.flatten().topk(k)
    return [unravel_index(i, x.shape) for i in indices]

