Say we have a matrix M that is n by n and after we flatten it we get n^2 sized array.
Then we find max at index k.
How to convert k to (x, y) coordinate for pre-flatten array?
Is there torch function for this?